WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-17 21:25:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

Remove duckdb checkpointer and store

Browse Source 
- duckdb is too buggy to be able to provide reliable checkpointer and store
This commit is contained in:
Nuno Campos
2025-01-15 09:15:17 -08:00
parent d12830e95c
commit 003b1883c9
24 changed files with 176 additions and 5076 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/ci.yml
+14 -6
View File
@@ -31,7 +31,6 @@ jobs:
"libs/cli",
"libs/checkpoint",
"libs/checkpoint-sqlite",
"libs/checkpoint-duckdb",
"libs/checkpoint-postgres",
"libs/scheduler-kafka",
]
@@ -44,12 +43,12 @@ jobs:
name: cd ${{ matrix.working-directory }}
strategy:
matrix:
working-directory: [
working-directory:
[
"libs/cli",
"libs/checkpoint",
"libs/checkpoint-sqlite",
"libs/checkpoint-duckdb",
"libs/checkpoint-postgres"
"libs/checkpoint-postgres",
]
uses: ./.github/workflows/_test.yml
with:
@@ -76,7 +75,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.11'
python-version: "3.11"
- name: Run check_sdk_methods script
run: python .github/scripts/check_sdk_methods.py
@@ -133,7 +132,16 @@ jobs:
ci_success:
name: "CI Success"
needs: [lint, lint-js, test, test-langgraph, test-scheduler-kafka, integration-test, test-js]
needs:
[
lint,
lint-js,
test,
test-langgraph,
test-scheduler-kafka,
integration-test,
test-js,
]
if: |
always()
runs-on: ubuntu-latest
libs/checkpoint-duckdb/Makefile
-35
View File
@@ -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
-95
View File
@@ -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
-356
View File
@@ -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 cls(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
-443
View File
@@ -1,443 +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 cls(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.
"""
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 `checkpointer.alist(...)` or `await "
"graph.ainvoke(...)`."
)
except RuntimeError:
pass
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
-290
View File
@@ -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/checkpoint/duckdb/py.typed
View File
libs/checkpoint-duckdb/langgraph/store/duckdb/__init__.py
-4
View File
@@ -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
-195
View File
@@ -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 cls(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
-408
View File
@@ -1,408 +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,
SearchItem,
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_search_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 _row_to_search_item(
namespace: tuple[str, ...],
row: tuple,
) -> SearchItem:
"""Convert a row from the database into an SearchItem."""
# TODO: Add support for search
_, key, val, created_at, updated_at = row
return SearchItem(
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/langgraph/store/duckdb/py.typed
View File
libs/checkpoint-duckdb/poetry.lock
Generated
-1058
View File
File diff suppressed because it is too large Load Diff
libs/checkpoint-duckdb/pyproject.toml
-60
View File
@@ -1,60 +0,0 @@
[tool.poetry]
name = "langgraph-checkpoint-duckdb"
version = "2.0.2"
description = "Library with a DuckDB implementation of LangGraph checkpoint saver."
authors = []
license = "MIT"
readme = "README.md"
repository = "https://www.github.com/langchain-ai/langgraph"
packages = [{ include = "langgraph" }]
[tool.poetry.dependencies]
python = "^3.9.0,<4.0"
langgraph-checkpoint = "^2.0.2"
duckdb = ">=1.1.2"
[tool.poetry.group.dev.dependencies]
ruff = "^0.6.2"
codespell = "^2.2.0"
pytest = "^7.2.1"
anyio = "^4.4.0"
pytest-asyncio = "^0.21.1"
pytest-mock = "^3.11.1"
pytest-watch = "^4.2.0"
mypy = "^1.10.0"
langgraph-checkpoint = {path = "../checkpoint", develop = true}
[tool.pytest.ini_options]
# --strict-markers will raise errors on unknown marks.
# https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
#
# https://docs.pytest.org/en/7.1.x/reference/reference.html
# --strict-config any warnings encountered while parsing the `pytest`
# section of the configuration file raise errors.
addopts = "--strict-markers --strict-config --durations=5 -vv"
asyncio_mode = "auto"
[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.ruff]
lint.select = [
"E", # pycodestyle
"F", # Pyflakes
"UP", # pyupgrade
"B", # flake8-bugbear
"I", # isort
]
lint.ignore = ["E501", "B008", "UP007", "UP006"]
[tool.mypy]
# https://mypy.readthedocs.io/en/stable/config_file.html
disallow_untyped_defs = "True"
explicit_package_bases = "True"
warn_no_return = "False"
warn_unused_ignores = "True"
warn_redundant_casts = "True"
allow_redefinition = "True"
disable_error_code = "typeddict-item, return-value"
libs/checkpoint-duckdb/tests/test_async.py
-112
View File
@@ -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
-517
View File
@@ -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
-457
View File
@@ -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
-111
View File
@@ -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/langgraph/poetry.lock
Generated
+162 -84
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/pyproject.toml
-1
View File
@@ -27,7 +27,6 @@ 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}
libs/langgraph/tests/__snapshots__/test_pregel.ambr
-647
View File
File diff suppressed because one or more lines are too long
libs/langgraph/tests/__snapshots__/test_pregel_async.ambr
-146
View File
@@ -329,127 +329,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;
@@ -1513,31 +1392,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
-39
View File
@@ -11,8 +11,6 @@ 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, ShallowPostgresSaver
from langgraph.checkpoint.postgres.aio import (
AsyncPostgresSaver,
@@ -21,7 +19,6 @@ from langgraph.checkpoint.postgres.aio import (
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
@@ -70,20 +67,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]}"
@@ -281,9 +264,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
@@ -370,13 +350,6 @@ async def _store_postgres_aio_pool():
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]}"
@@ -433,13 +406,6 @@ def store_postgres_pool():
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()
@@ -460,9 +426,6 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
elif store_name == "postgres_aio_pool":
async with _store_postgres_aio_pool() 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}")
@@ -500,12 +463,10 @@ ALL_STORES_SYNC = [
"postgres",
"postgres_pipe",
"postgres_pool",
"duckdb",
]
ALL_STORES_ASYNC = [
"in_memory",
"postgres_aio",
"postgres_aio_pipe",
"postgres_aio_pool",
"duckdb_aio",
]
libs/langgraph/tests/test_pregel.py
-6
View File
@@ -4159,12 +4159,6 @@ def test_store_injected(
builder.add_edge("__start__", "node")
N = 500
M = 1
if "duckdb" in store_name:
logger.warning(
"DuckDB store implementation has a known issue that does not"
" support concurrent writes, so we're reducing the test scope"
)
N = M = 1
for i in range(N):
builder.add_node(f"node_{i}", Node(i))
libs/langgraph/tests/test_pregel_async.py
-6
View File
@@ -5750,12 +5750,6 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) ->
N = 500
M = 1
if "duckdb" in store_name:
logger.warning(
"DuckDB store implementation has a known issue that does not"
" support concurrent writes, so we're reducing the test scope"
)
N = M = 1
for i in range(N):
builder.add_node(f"node_{i}", Node(i))