x

.github/workflows/_test.yml

@@ -35,7 +35,6 @@ 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,7 +37,6 @@ 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,7 +29,6 @@ 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,7 +31,6 @@ jobs:
             "libs/cli",
             "libs/checkpoint",
             "libs/checkpoint-sqlite",
-            "libs/checkpoint-duckdb",
             "libs/checkpoint-postgres",
             "libs/scheduler-kafka",
           ]
@@ -48,7 +47,6 @@ 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,7 +81,6 @@ 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 use. This is parameterized by a user defined `State` object.
+The `StateGraph` class is the main graph class to uses. 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, "a-memory", {"rules": ["User likes short, direct language", "User only speaks English & python"], "my-key": "my-value"})
+store.put(namespace, key="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, "a-memory")
+item = store.get(namespace)
 # 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](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.
+- **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.
 - **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,20 +47,13 @@ from langgraph.graph import StateGraph, MessagesState, START
 model = ChatOpenAI()
 
 class AgentState(MessagesState):
-    next: Literal["agent_1", "agent_2", "__end__"]
+    next: Literal["agent_1", "agent_2"]
 
 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]}
 
@@ -75,7 +68,6 @@ 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")
@@ -96,117 +88,18 @@ 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]):
-    # 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
+    tool_message = ...
+    return {"messages": [tool_message]}
 
 def agent_2(state: Annotated[dict, InjectedState]):
-    response = model.invoke(...)
-    return response.content
+    tool_message = ...
+    return {"messages": [tool_message]}
 
 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:

libs/checkpoint-duckdb/Makefile

@@ -1,35 +0,0 @@
-.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

@@ -1,95 +0,0 @@
-# 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

@@ -1,356 +0,0 @@
-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

@@ -1,431 +0,0 @@
-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

@@ -1,290 +0,0 @@
-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

@@ -1,4 +0,0 @@
-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

@@ -1,195 +0,0 @@
-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

@@ -1,391 +0,0 @@
-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

@@ -1,60 +0,0 @@
-[tool.poetry]
-name = "langgraph-checkpoint-duckdb"
-version = "2.0.0"
-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.0"
-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

@@ -1,112 +0,0 @@
-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

@@ -1,517 +0,0 @@
-# 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

@@ -1,457 +0,0 @@
-# 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

@@ -1,111 +0,0 @@
-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", "channel_values": {"key": "value"}}
+            >>>     checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "data": {"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, before.
+        """Return WHERE clause predicates for alist() given config, filter, cursor.
 
         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,6 +44,7 @@ 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-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", "channel_values": {"key": "value"}}
+            >>>     checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "data": {"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/langgraph/langgraph/graph/state.py

@@ -3,7 +3,7 @@ import logging
 import typing
 import warnings
 from functools import partial
-from inspect import isclass, isfunction, ismethod, signature
+from inspect import isclass, isfunction, signature
 from typing import (
     Any,
     Callable,
@@ -338,8 +338,19 @@ class StateGraph(Graph):
                     f"'{character}' is a reserved character and is not allowed in the node names."
                 )
 
-        if input is None:
-            input = _get_input_schema_from_type_hint(action)
+        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 not None:
             self._add_schema(input)
         self.nodes[cast(str, node)] = StateNodeSpec(
@@ -823,21 +834,3 @@ 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

@@ -70,8 +70,8 @@ def msg_content_output(output: Any) -> str | List[dict]:
 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 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
+    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
     a list of ToolMessages, one for each tool call.
 
     The `ToolNode` is roughly analogous to:
@@ -102,14 +102,12 @@ class ToolNode(RunnableCallable):
         name: str = "tools",
         tags: Optional[list[str]] = None,
         handle_tool_errors: Optional[bool] = 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_))
@@ -133,7 +131,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 {self.messages_key: outputs}
+        return outputs if output_type == "list" else {"messages": outputs}
 
     def invoke(
         self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
@@ -165,7 +163,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 {self.messages_key: outputs}
+        return outputs if output_type == "list" else {"messages": outputs}
 
     def _run_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
         if invalid_tool_message := self._validate_tool_call(call):
@@ -216,10 +214,10 @@ class ToolNode(RunnableCallable):
         if isinstance(input, list):
             output_type = "list"
             message: AnyMessage = input[-1]
-        elif isinstance(input, dict) and (messages := input.get(self.messages_key, [])):
+        elif isinstance(input, dict) and (messages := input.get("messages", [])):
             output_type = "dict"
             message = messages[-1]
-        elif messages := getattr(input, self.messages_key, None):
+        elif messages := getattr(input, "messages", None):
             # Assume dataclass-like state that can coerce from dict
             output_type = "dict"
             message = messages[-1]
@@ -258,10 +256,10 @@ class ToolNode(RunnableCallable):
             required_fields = list(state_args.values())
             if (
                 len(required_fields) == 1
-                and required_fields[0] == self.messages_key
+                and required_fields[0] == "messages"
                 or required_fields[0] is None
             ):
-                input = {self.messages_key: input}
+                input = {"messages": input}
             else:
                 err_msg = (
                     f"Invalid input to ToolNode. Tool {tool_call['name']} requires "
@@ -327,7 +325,6 @@ 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
 
@@ -380,9 +377,9 @@ def tools_condition(
     """
     if isinstance(state, list):
         ai_message = state[-1]
-    elif isinstance(state, dict) and (messages := state.get(messages_key, [])):
+    elif isinstance(state, dict) and (messages := state.get("messages", [])):
         ai_message = messages[-1]
-    elif messages := getattr(state, messages_key, []):
+    elif messages := getattr(state, "messages", []):
         ai_message = messages[-1]
     else:
         raise ValueError(f"No messages found in input state to tool_edge: {state}")

libs/langgraph/langgraph/pregel/executor.py

@@ -118,17 +118,11 @@ class AsyncBackgroundExecutor(AsyncContextManager):
     - re-raises the first exception from tasks with `__reraise_on_exit__=True`
       ignoring CancelledError"""
 
-    def __init__(self, config: RunnableConfig) -> None:
+    def __init__(self) -> 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,
@@ -140,8 +134,6 @@ 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:
@@ -191,9 +183,3 @@ 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,9 +847,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             else []
         )
 
-        self.submit = await self.stack.enter_async_context(
-            AsyncBackgroundExecutor(self.config)
-        )
+        self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor())
         self.channels, self.managed = await self.stack.enter_async_context(
             AsyncChannelsManager(self.specs, self.checkpoint, self)
         )

libs/langgraph/langgraph/pregel/remote.py

@@ -19,6 +19,7 @@ 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,
@@ -26,34 +27,25 @@ 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 RemoteException(Exception):
-    """Exception raised when an error occurs in the remote graph."""
-
-    pass
-
-
 class RemoteGraph(PregelProtocol, Runnable):
     def __init__(
         self,
         graph_id: str,
+        config: Optional[RunnableConfig] = None,
         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.
 
@@ -356,37 +348,6 @@ 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],
@@ -399,40 +360,18 @@ 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=cast(str, sanitized_config["configurable"]["thread_id"]),
+            thread_id=sanitized_config["configurable"]["thread_id"],
             assistant_id=self.graph_id,
             input=input,
             config=sanitized_config,
-            stream_mode=stream_modes,
-            interrupt_before=interrupt_before,
-            interrupt_after=interrupt_after,
+            stream_mode=stream_mode,  # type: ignore
+            interrupt_before=interrupt_before,  # type: ignore
+            interrupt_after=interrupt_after,  # type: ignore
             stream_subgraphs=subgraphs,
-            if_not_exists="create",
         ):
-            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
+            yield chunk
 
     async def astream(
         self,
@@ -446,40 +385,47 @@ 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,
             input=input,
             config=sanitized_config,
-            stream_mode=stream_modes,
-            interrupt_before=interrupt_before,
-            interrupt_after=interrupt_after,
+            stream_mode=stream_mode if stream_mode else "values",  # type: ignore
+            interrupt_before=interrupt_before,  # type: ignore
+            interrupt_after=interrupt_after,  # type: ignore
             stream_subgraphs=subgraphs,
-            if_not_exists="create",
         ):
-            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
+            yield chunk
+
+    async def astream_events(
+        self,
+        input: Any,
+        config: Optional[RunnableConfig] = 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,
+            )
 
     def invoke(
         self,
@@ -497,9 +443,8 @@ class RemoteGraph(PregelProtocol, Runnable):
             assistant_id=self.graph_id,
             input=input,
             config=sanitized_config,
-            interrupt_before=interrupt_before,
-            interrupt_after=interrupt_after,
-            if_not_exists="create",
+            interrupt_before=interrupt_before,  # type: ignore
+            interrupt_after=interrupt_after,  # type: ignore
         )
 
     async def ainvoke(
@@ -518,7 +463,6 @@ class RemoteGraph(PregelProtocol, Runnable):
             assistant_id=self.graph_id,
             input=input,
             config=sanitized_config,
-            interrupt_before=interrupt_before,
-            interrupt_after=interrupt_after,
-            if_not_exists="create",
+            interrupt_before=interrupt_before,  # type: ignore
+            interrupt_after=interrupt_after,  # type: ignore
         )

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.39"
+version = "0.2.38"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"
@@ -27,12 +27,11 @@ 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.0beta1"
+uvloop = "0.21.0"
 pyperf = "^2.7.0"
 py-spy = "^0.3.14"
 types-requests = "^2.32.0.20240914"

libs/langgraph/tests/__snapshots__/test_pregel_async.ambr

@@ -221,6 +221,19 @@
                 +---------+                
   '''
 # ---
+# 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({
@@ -329,127 +342,6 @@
     '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;
@@ -1302,31 +1194,6 @@
    +---------+   
   '''
 # ---
-# 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,14 +11,11 @@ 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
@@ -64,20 +61,6 @@ 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]}"
@@ -229,9 +212,6 @@ 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
@@ -267,13 +247,6 @@ 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]}"
@@ -291,13 +264,6 @@ 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()
@@ -312,9 +278,6 @@ 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}")
 
@@ -322,7 +285,6 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
 ALL_CHECKPOINTERS_SYNC = [
     "memory",
     "sqlite",
-    "duckdb",
     "postgres",
     "postgres_pipe",
     "postgres_pool",
@@ -330,7 +292,6 @@ ALL_CHECKPOINTERS_SYNC = [
 ALL_CHECKPOINTERS_ASYNC = [
     "memory",
     "sqlite_aio",
-    "duckdb_aio",
     "postgres_aio",
     "postgres_aio_pipe",
     "postgres_aio_pool",
@@ -339,5 +300,5 @@ ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [
     *ALL_CHECKPOINTERS_ASYNC,
     None,
 ]
-ALL_STORES_SYNC = ["in_memory", "postgres", "duckdb"]
-ALL_STORES_ASYNC = ["in_memory", "postgres_aio", "duckdb_aio"]
+ALL_STORES_SYNC = ["in_memory", "postgres"]
+ALL_STORES_ASYNC = ["in_memory", "postgres_aio"]

libs/langgraph/tests/test_prebuilt.py

@@ -1,6 +1,5 @@
 import dataclasses
 import json
-from functools import partial
 from typing import (
     Annotated,
     Any,
@@ -36,13 +35,8 @@ from pydantic.v1 import BaseModel as BaseModelV1
 from typing_extensions import TypedDict
 
 from langgraph.checkpoint.base import BaseCheckpointSaver
-from langgraph.graph import START, MessagesState, StateGraph, add_messages
-from langgraph.prebuilt import (
-    ToolNode,
-    ValidationNode,
-    create_react_agent,
-    tools_condition,
-)
+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.store.base import BaseStore
 from langgraph.store.memory import InMemoryStore
@@ -52,7 +46,7 @@ from tests.conftest import (
     IS_LANGCHAIN_CORE_030_OR_GREATER,
     awith_checkpointer,
 )
-from tests.messages import _AnyIdHumanMessage, _AnyIdToolMessage
+from tests.messages import _AnyIdHumanMessage
 
 pytestmark = pytest.mark.anyio
 
@@ -832,47 +826,6 @@ 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,53 +1922,6 @@ 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_pregel.py

@@ -7,9 +7,7 @@ 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
 
@@ -475,97 +473,17 @@ 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(sync_client=mock_sync_client, graph_id="test_graph_id")
-
-    # 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"},
     ]
 
-    mock_sync_client.runs.stream.return_value = [
-        StreamPart(event="updates", data={"chunk": "data3"}),
-        StreamPart(event="updates", data={"chunk": "data4"}),
-        StreamPart(event="updates", data={"__interrupt__": ()}),
-    ]
+    # call method / assertions
+    remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
 
-    # 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"}),
-    ]
+    config = {"configurable": {"thread_id": "thread_1"}}
+    result = list(remote_pregel.stream({"input": "data"}, config))
+    assert result == [{"chunk": "data1"}, {"chunk": "data2"}, {"chunk": "data3"}]
 
 
 @pytest.mark.anyio
@@ -574,139 +492,20 @@ async def test_astream():
     mock_async_client = MagicMock()
     async_iter = MagicMock()
     async_iter.__aiter__.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__": ()}),
+        {"chunk": "data1"},
+        {"chunk": "data2"},
+        {"chunk": "data3"},
     ]
     mock_async_client.runs.stream.return_value = async_iter
 
     # call method / assertions
     remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
 
-    # 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"}),
-    ]
+    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"}]
 
 
 def test_invoke():
@@ -773,7 +572,7 @@ async def test_langgraph_cloud_integration():
         "messages": [
             {
                 "role": "human",
-                "content": "What's the weather in SF?",
+                "content": "Hello world!",
             }
         ]
     }
@@ -781,8 +580,7 @@ async def test_langgraph_cloud_integration():
     # test invoke
     response = app.invoke(
         input,
-        config={"configurable": {"thread_id": "39a6104a-34e7-4f83-929c-d9eb163003c9"}},
-        interrupt_before=["agent"],
+        config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
     )
     print("response:", response["messages"][-1].content)
 

libs/langgraph/tests/test_state.py

@@ -61,9 +61,6 @@ 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}
 
@@ -76,46 +73,24 @@ 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,
-        pre_foo,
-        Foo(),
-    ]
+    actions = [complete_hint, miss_first_hint, only_return_hint, miss_all_hint]
 
     for action in actions:
         graph.add_node(action)
 
-    def get_name(action) -> str:
-        return getattr(action, "__name__", action.__class__.__name__)
-
-    graph.set_entry_point(get_name(actions[0]))
+    graph.set_entry_point(actions[0].__name__)
     for i in range(len(actions) - 1):
-        graph.add_edge(get_name(actions[i]), get_name(actions[i + 1]))
-    graph.set_finish_point(get_name(actions[-1]))
+        graph.add_edge(actions[i].__name__, actions[i + 1].__name__)
+    graph.set_finish_point(actions[-1].__name__)
 
     graph = graph.compile()
 
     input_state = InputState(question="Hello World!")
     output_state = OutputState(input_state=input_state)
-    foo_state = FooState(foo="bar")
     for i, c in enumerate(graph.stream(input_state, stream_mode="updates")):
-        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
+        node_name = actions[i].__name__
+        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(msg["config"]) as submit:
+        ) as (channels, managed), AsyncBackgroundExecutor() as submit:
             if task := await asyncio.to_thread(
                 prepare_single_task,
                 msg["task"]["path"],

libs/sdk-py/langgraph_sdk/client.py

@@ -39,7 +39,6 @@ from langgraph_sdk.schema import (
     Cron,
     DisconnectMode,
     GraphSchema,
-    IfNotExists,
     Item,
     Json,
     ListNamespaceResponse,
@@ -1170,19 +1169,18 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[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, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
-        feedback_keys: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
+        feedback_keys: Optional[list[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]: ...
 
@@ -1193,13 +1191,13 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
-        feedback_keys: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
+        feedback_keys: Optional[list[str]] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
         webhook: Optional[str] = None,
@@ -1212,20 +1210,19 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[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, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
-        feedback_keys: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
+        feedback_keys: Optional[list[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.
@@ -1251,8 +1248,6 @@ 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.
 
@@ -1298,7 +1293,6 @@ 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,
@@ -1319,12 +1313,12 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
         webhook: Optional[str] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
         after_seconds: Optional[int] = None,
@@ -1337,17 +1331,16 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[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, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
-        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Run: ...
 
@@ -1357,17 +1350,16 @@ class RunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[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, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[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:
@@ -1391,8 +1383,6 @@ 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.
 
@@ -1476,7 +1466,6 @@ 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,
         }
@@ -1506,12 +1495,11 @@ class RunsClient:
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[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]]: ...
 
@@ -1524,8 +1512,8 @@ class RunsClient:
         input: Optional[dict] = None,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
         webhook: Optional[str] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
@@ -1542,13 +1530,12 @@ class RunsClient:
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[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.
@@ -1571,8 +1558,6 @@ 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.
 
@@ -1634,7 +1619,6 @@ 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,
@@ -3248,19 +3232,18 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[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, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
-        feedback_keys: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
+        feedback_keys: Optional[list[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]: ...
 
@@ -3271,13 +3254,13 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
-        feedback_keys: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
+        feedback_keys: Optional[list[str]] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
         webhook: Optional[str] = None,
@@ -3290,20 +3273,19 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[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, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
-        feedback_keys: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
+        feedback_keys: Optional[list[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.
@@ -3329,8 +3311,6 @@ 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.
 
@@ -3376,7 +3356,6 @@ 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,
@@ -3397,12 +3376,12 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[StreamMode]] = "values",
         stream_subgraphs: bool = False,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
         webhook: Optional[str] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
         after_seconds: Optional[int] = None,
@@ -3415,17 +3394,16 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[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, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
-        if_not_exists: Optional[IfNotExists] = None,
         after_seconds: Optional[int] = None,
     ) -> Run: ...
 
@@ -3435,18 +3413,17 @@ class SyncRunsClient:
         assistant_id: str,
         *,
         input: Optional[dict] = None,
-        stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
+        stream_mode: Union[StreamMode, list[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, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[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.
@@ -3469,8 +3446,6 @@ 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.
 
@@ -3554,7 +3529,6 @@ 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,
         }
@@ -3584,12 +3558,11 @@ class SyncRunsClient:
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[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]]: ...
 
@@ -3602,8 +3575,8 @@ class SyncRunsClient:
         input: Optional[dict] = None,
         metadata: Optional[dict] = None,
         config: Optional[Config] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[str]]] = None,
         webhook: Optional[str] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         on_completion: Optional[OnCompletionBehavior] = None,
@@ -3620,13 +3593,12 @@ class SyncRunsClient:
         config: Optional[Config] = None,
         checkpoint: Optional[Checkpoint] = None,
         checkpoint_id: Optional[str] = None,
-        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
-        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_before: Optional[Union[All, list[str]]] = None,
+        interrupt_after: Optional[Union[All, list[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.
@@ -3649,8 +3621,6 @@ 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.
 
@@ -3712,7 +3682,6 @@ 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

@@ -69,13 +69,6 @@ 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."""