Compare commits

..
23 Commits
Author SHA1 Message Date
Nuno CamposandGitHub 2e1adaa867 0.3.12 2025-03-17 16:56:34 -07:00
Nuno CamposandGitHub 3f8b165592 Update state.py 2025-03-17 16:56:09 -07:00
William FHandGitHub 9ed0fa196c langgraph-checkpoint 2.0.21 (#3883) 2025-03-17 15:25:26 -07:00
William Fu-Hinthorn 0b9adc28c3 langgraph-checkpoint 2.0.21 2025-03-17 15:15:29 -07:00
William FHandGitHub 3b0255d1ef Check that migrations are idempotent (#3881) 2025-03-17 15:13:59 -07:00
William FHandGitHub d4255a0645 Merge branch 'main' into wfh/idempotency_test_ 2025-03-17 13:27:12 -07:00
William FHandGitHub 80d61a2600 Make expires_at idempotent (#3880) 2025-03-17 13:26:20 -07:00
William Fu-Hinthorn 424f24720a Make expires_at idempotent 2025-03-17 13:25:22 -07:00
William Fu-Hinthorn 2a71180c1d Add tests for idempotency in migraionts 2025-03-17 12:43:21 -07:00
William Fu-Hinthorn 697f878e36 Make expires_at idempotent 2025-03-17 12:38:43 -07:00
alxdr3kandGitHub def69c59d2 docs: add missing namespace definition in persistence.md (#3859)
- fix typo
- add missing namespace definition for memory store
2025-03-17 09:38:15 -04:00
Nuno Campos bad4d17c34 Remove 'check size' ci job 2025-03-14 21:54:01 -07:00
Nuno Campos 55219b23d8 0.3.11 2025-03-14 16:18:50 -07:00
Nuno CamposandGitHub 8edbd39ad3 Add optional encryption of checkpointer payloads (#3852)
- no dependency on any particular encryption lib (there is no py stdlib
encryption lib)
- works with any modern checkpointer, ie. those which use dumps_typed
and loads_typed methods to serialize data
- uses the default msg pack serializer, but also works with any custom
serializer
- backwards compatible with unencrypted data in same storage (will just
be read unencrypted)
- providing easy constructor to use AES encryption through pycriptodome
library, one single line of code to add it in
- other encryption libraries or algorithms (even assymetric ones) can be
used by implementing the two-method CipherProtocol interface
- cipher name (eg. aes) is stored with encrypted payload for forwards
compatibility

```py
import sqlite3

from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
from langgraph.checkpoint.sqlite import SqliteSaver

# will read AES key from env var LANGGRAPH_AES_KEY
serde = EncryptedSerializer.from_pycryptodome_aes()
# works with any other checkpointer, including custom ones
checkpointer = SqliteSaver(sqlite3.connect('...'), serde=serde)
```
2025-03-14 16:14:29 -07:00
Nuno Campos 4b0fd834d8 Make it possible to implement a cipher that handles multiple protocols 2025-03-14 16:04:28 -07:00
Nuno Campos 0fd2748530 Accept custom serde implementations 2025-03-14 15:43:07 -07:00
Nuno Campos bc0a3419ed Lint 2025-03-14 15:38:22 -07:00
Nuno Campos 5cd47bac49 Add optional encryption of checkpointer payloads
- no dependency on any particular encryption lib (there is no py stdlib encryption lib)
- works with any modern checkpointer, ie. those which use dumps_typed and loads_typed methods to serialize data
- backwards compatible with unencrypted data in same storage (will just be read unencrypted)
- providing easy constructor to use AES encryption through pycriptodome library, one single line of code to add it in
- other encryption libraries or algorithms (even assymetric ones) can be used by implementing the two-method CipherProtocol interface
- cipher name (eg. aes) is stored with encrypted payload for forwards compatibility
2025-03-14 15:32:50 -07:00
William FHandGitHub 4c6d80a67f Add TTL Sweeper (#3849) 2025-03-14 14:29:40 -07:00
William Fu-Hinthorn 18ed044c27 Bump patch version 2025-03-14 14:03:30 -07:00
William Fu-Hinthorn 394a9fa85f Update schema 2025-03-14 13:48:23 -07:00
William Fu-Hinthorn 9741d9bdf0 Add tests for sweeper (sync) 2025-03-14 13:43:53 -07:00
William Fu-Hinthorn 06ca07432d Add sweeper 2025-03-13 18:55:16 -07:00
22 changed files with 1587 additions and 257 deletions
-29
View File
@@ -1,29 +0,0 @@
name: Check File Size
on:
push:
branches:
- main
pull_request:
branches:
- main
workflow_dispatch:
jobs:
file-size-check:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Get changed files
id: changed-files
uses: tj-actions/changed-files@v44
- name: Filter by size
# TODO: roll back the web voyager hack
run: |
large_added_files=$(find ${{ steps.changed-files.outputs.added_files }} -maxdepth 0 -size +1M | grep -v "web_voyager" || true)
if [ -n "$large_added_files" ]; then
echo "Large files added: $large_added_files"
echo "# Large files added:" >> $GITHUB_STEP_SUMMARY
echo "$large_added_files" >> $GITHUB_STEP_SUMMARY
exit 1
fi
+4 -1
View File
@@ -232,7 +232,7 @@ from langgraph.store.memory import InMemoryStore
in_memory_store = InMemoryStore()
```
Memories are namespaced by a `tuple`, which in this specific example will be `(<user_id>, "memories")`. The namespace can be any length and represent anything, does not have be user specific.
Memories are namespaced by a `tuple`, which in this specific example will be `(<user_id>, "memories")`. The namespace can be any length and represent anything, does not have to be user specific.
```python
user_id = "1"
@@ -387,6 +387,9 @@ We can access the memories and use them in our model call.
def call_model(state: MessagesState, config: RunnableConfig, *, store: BaseStore):
# Get the user id from the config
user_id = config["configurable"]["user_id"]
# Namespace the memory
namespace = (user_id, "memories")
# Search based on the most recent message
memories = store.search(
@@ -2,6 +2,7 @@ import asyncio
import logging
from collections.abc import AsyncIterator, Iterable, Sequence
from contextlib import asynccontextmanager
from types import TracebackType
from typing import Any, Callable, Optional, Union, cast
import orjson
@@ -25,6 +26,7 @@ from langgraph.store.postgres.base import (
PoolConfig,
PostgresIndexConfig,
Row,
TTLConfig,
_decode_ns_bytes,
_ensure_index_config,
_group_ops,
@@ -106,6 +108,11 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
Semantic search is disabled by default. You can enable it by providing an `index` configuration
when creating the store. Without this configuration, all `index` arguments passed to
`put` or `aput` will have no effect.
Note:
If you provide a TTL configuration, you must explicitly call `start_ttl_sweeper()` to begin
the background task that removes expired items. Call `stop_ttl_sweeper()` to properly
clean up resources when you're done with the store.
"""
__slots__ = (
@@ -115,7 +122,11 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
"supports_pipeline",
"index_config",
"embeddings",
"ttl_config",
"_ttl_sweeper_task",
"_ttl_stop_event",
)
supports_ttl: bool = True
def __init__(
self,
@@ -126,6 +137,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
] = None,
index: Optional[PostgresIndexConfig] = None,
ttl: Optional[TTLConfig] = None,
) -> None:
if isinstance(conn, AsyncConnectionPool) and pipe is not None:
raise ValueError(
@@ -141,10 +153,13 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
self.index_config = index
if self.index_config:
self.embeddings, self.index_config = _ensure_index_config(self.index_config)
else:
self.embeddings = None
self.ttl_config = ttl
self._ttl_sweeper_task: Optional[asyncio.Task[None]] = None
self._ttl_stop_event = asyncio.Event()
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
grouped_ops, num_ops = _group_ops(ops)
results: list[Result] = [None] * num_ops
@@ -167,6 +182,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
pipeline: bool = False,
pool_config: Optional[PoolConfig] = None,
index: Optional[PostgresIndexConfig] = None,
ttl: Optional[TTLConfig] = None,
) -> AsyncIterator["AsyncPostgresStore"]:
"""Create a new AsyncPostgresStore instance from a connection string.
@@ -198,16 +214,16 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
**cast(dict, pc),
),
) as pool:
yield cls(conn=pool, index=index)
yield cls(conn=pool, index=index, ttl=ttl)
else:
async with await AsyncConnection.connect(
conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
) as conn:
if pipeline:
async with conn.pipeline() as pipe:
yield cls(conn=conn, pipe=pipe, index=index)
yield cls(conn=conn, pipe=pipe, index=index, ttl=ttl)
else:
yield cls(conn=conn, index=index)
yield cls(conn=conn, index=index, ttl=ttl)
async def setup(self) -> None:
"""Set up the store database asynchronously.
@@ -256,6 +272,119 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
"INSERT INTO vector_migrations (v) VALUES (%s)", (v,)
)
async def sweep_ttl(self) -> int:
"""Delete expired store items based on TTL.
Returns:
int: The number of deleted items.
"""
async with self._cursor() as cur:
await cur.execute(
"""
DELETE FROM store
WHERE expires_at IS NOT NULL AND expires_at < NOW()
"""
)
deleted_count = cur.rowcount
return deleted_count
async def start_ttl_sweeper(
self, sweep_interval_minutes: Optional[int] = None
) -> asyncio.Task[None]:
"""Periodically delete expired store items based on TTL.
Returns:
Task that can be awaited or cancelled.
"""
if not self.ttl_config:
return asyncio.create_task(asyncio.sleep(0))
if self._ttl_sweeper_task is not None and not self._ttl_sweeper_task.done():
return self._ttl_sweeper_task
self._ttl_stop_event.clear()
interval = float(
sweep_interval_minutes or self.ttl_config.get("sweep_interval_minutes") or 5
)
logger.info(f"Starting store TTL sweeper with interval {interval} minutes")
async def _sweep_loop() -> None:
while not self._ttl_stop_event.is_set():
try:
try:
await asyncio.wait_for(
self._ttl_stop_event.wait(),
timeout=interval * 60,
)
break
except asyncio.TimeoutError:
pass
expired_items = await self.sweep_ttl()
if expired_items > 0:
logger.info(f"Store swept {expired_items} expired items")
except asyncio.CancelledError:
break
except Exception as exc:
logger.exception("Store TTL sweep iteration failed", exc_info=exc)
task = asyncio.create_task(_sweep_loop())
task.set_name("ttl_sweeper")
self._ttl_sweeper_task = task
return task
async def stop_ttl_sweeper(self, timeout: Optional[float] = None) -> bool:
"""Stop the TTL sweeper task if it's running.
Args:
timeout: Maximum time to wait for the task to stop, in seconds.
If None, wait indefinitely.
Returns:
bool: True if the task was successfully stopped or wasn't running,
False if the timeout was reached before the task stopped.
"""
if self._ttl_sweeper_task is None or self._ttl_sweeper_task.done():
return True
logger.info("Stopping TTL sweeper task")
self._ttl_stop_event.set()
if timeout is not None:
try:
await asyncio.wait_for(self._ttl_sweeper_task, timeout=timeout)
success = True
except asyncio.TimeoutError:
success = False
else:
await self._ttl_sweeper_task
success = True
if success:
self._ttl_sweeper_task = None
logger.info("TTL sweeper task stopped")
else:
logger.warning("Timed out waiting for TTL sweeper task to stop")
return success
async def __aenter__(self) -> "AsyncPostgresStore":
return self
async def __aexit__(
self,
exc_type: Optional[type[BaseException]],
exc_val: Optional[BaseException],
exc_tb: Optional["TracebackType"],
) -> None:
# Ensure the TTL sweeper task is stopped when exiting the context
if hasattr(self, "_ttl_sweeper_task") and self._ttl_sweeper_task is not None:
# Set the event to signal the task to stop
self._ttl_stop_event.set()
# We don't wait for the task to complete here to avoid blocking
# The task will clean up itself gracefully
async def _execute_batch(
self,
grouped_ops: dict,
@@ -1,4 +1,5 @@
import asyncio
import concurrent.futures
import json
import logging
import threading
@@ -74,6 +75,17 @@ CREATE TABLE IF NOT EXISTS store (
"""
-- For faster lookups by prefix
CREATE INDEX CONCURRENTLY IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
""",
"""
-- Add expires_at column to store table
ALTER TABLE store
ADD COLUMN IF NOT EXISTS expires_at TIMESTAMP WITH TIME ZONE,
ADD COLUMN IF NOT EXISTS ttl_minutes INT;
""",
"""
-- Add indexes for efficient TTL sweeping
CREATE INDEX IF NOT EXISTS idx_store_expires_at ON store (expires_at)
WHERE expires_at IS NOT NULL;
""",
]
@@ -225,20 +237,55 @@ class BasePostgresStore(Generic[C]):
self,
get_ops: Sequence[tuple[int, GetOp]],
) -> list[tuple[str, tuple, tuple[str, ...], list]]:
"""
Build queries to fetch (and optionally refresh the TTL of) multiple keys per namespace.
Each returned element is a tuple of:
(sql_query_string, sql_params, namespace, items_for_this_namespace)
where items_for_this_namespace is the original list of (idx, key, refresh_ttl).
"""
namespace_groups = defaultdict(list)
refresh_ttls = defaultdict(list)
for idx, op in get_ops:
namespace_groups[op.namespace].append((idx, op.key))
refresh_ttls[op.namespace].append(op.refresh_ttl)
results = []
for namespace, items in namespace_groups.items():
_, keys = zip(*items)
keys_to_query = ",".join(["%s"] * len(keys))
query = f"""
SELECT key, value, created_at, updated_at
FROM store
WHERE prefix = %s AND key IN ({keys_to_query})
this_refresh_ttls = refresh_ttls[namespace]
query = """
WITH passed_in AS (
SELECT unnest(%s::text[]) AS key,
unnest(%s::bool[]) AS do_refresh
),
updated AS (
UPDATE store s
SET expires_at = NOW() + (s.ttl_minutes || ' minutes')::interval
FROM passed_in p
WHERE s.prefix = %s
AND s.key = p.key
AND p.do_refresh = TRUE
AND s.ttl_minutes IS NOT NULL
RETURNING s.key
)
SELECT s.key, s.value, s.created_at, s.updated_at
FROM store s
JOIN passed_in p ON s.key = p.key
WHERE s.prefix = %s
"""
params = (_namespace_to_text(namespace), *keys)
ns_text = _namespace_to_text(namespace)
params = (
list(keys), # -> unnest(%s::text[])
list(this_refresh_ttls), # -> unnest(%s::bool[])
ns_text, # -> prefix = %s (for UPDATE)
ns_text, # -> prefix = %s (for final SELECT)
)
results.append((query, params, namespace, items))
return results
def _prepare_batch_PUT_queries(
@@ -248,7 +295,6 @@ class BasePostgresStore(Generic[C]):
list[tuple[str, Sequence]],
Optional[tuple[str, Sequence[tuple[str, str, str, str]]]],
]:
# Last-write wins
dedupped_ops: dict[tuple[tuple[str, ...], str], PutOp] = {}
for _, op in put_ops:
dedupped_ops[(op.namespace, op.key)] = op
@@ -282,15 +328,26 @@ class BasePostgresStore(Generic[C]):
insertion_params = []
vector_values = []
embedding_request_params = []
# Handle TTL expiration
# First handle main store insertions
for op in inserts:
values.append("(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)")
if op.ttl is not None:
expires_at_str = f"NOW() + INTERVAL '{op.ttl*60} seconds'"
ttl_minutes = op.ttl
else:
expires_at_str = "NULL"
ttl_minutes = None
values.append(
f"(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, {expires_at_str}, %s)"
)
insertion_params.extend(
[
_namespace_to_text(op.namespace),
op.key,
Jsonb(cast(dict, op.value)),
ttl_minutes,
]
)
@@ -304,7 +361,7 @@ class BasePostgresStore(Generic[C]):
k = op.key
if op.index is None:
paths = self.index_config["__tokenized_fields"]
paths = cast(dict, self.index_config)["__tokenized_fields"]
else:
paths = [(ix, tokenize_path(ix)) for ix in op.index]
@@ -319,11 +376,13 @@ class BasePostgresStore(Generic[C]):
values_str = ",".join(values)
query = f"""
INSERT INTO store (prefix, key, value, created_at, updated_at)
INSERT INTO store (prefix, key, value, created_at, updated_at, expires_at, ttl_minutes)
VALUES {values_str}
ON CONFLICT (prefix, key) DO UPDATE
SET value = EXCLUDED.value,
updated_at = CURRENT_TIMESTAMP
updated_at = CURRENT_TIMESTAMP,
expires_at = EXCLUDED.expires_at,
ttl_minutes = EXCLUDED.ttl_minutes
"""
queries.append((query, insertion_params))
@@ -347,92 +406,105 @@ class BasePostgresStore(Generic[C]):
list[tuple[str, list[Union[None, str, list[float]]]]], # queries, params
list[tuple[int, str]], # idx, query_text pairs to embed
]:
"""
Build per-SearchOp SQL queries (with optional TTL refresh) plus embedding requests.
Returns:
- queries: list of (SQL, param_list)
- embedding_requests: list of (original_index_in_search_ops, text_query)
"""
queries = []
embedding_requests = []
for idx, (_, op) in enumerate(search_ops):
# Build filter conditions first
filter_params = []
filter_conditions = []
filter_clauses = []
if op.filter:
for key, value in op.filter.items():
if isinstance(value, dict):
for op_name, val in value.items():
condition, filter_params_ = self._get_filter_condition(
condition, params_ = self._get_filter_condition(
key, op_name, val
)
filter_conditions.append(condition)
filter_params.extend(filter_params_)
filter_clauses.append(condition)
filter_params.extend(params_)
else:
filter_conditions.append("value->%s = %s::jsonb")
filter_params.extend([key, json.dumps(value)])
filter_clauses.append("value->%s = %s::jsonb")
filter_params.extend([key, orjson.dumps(value).decode("utf-8")])
ns_condition = "TRUE"
ns_param: Optional[Sequence[Union[str]]] = None
if op.namespace_prefix:
ns_condition = "store.prefix LIKE %s"
ns_param = (f"{_namespace_to_text(op.namespace_prefix)}%",)
else:
ns_param = ()
extra_filters = (
" AND " + " AND ".join(filter_clauses) if filter_clauses else ""
)
# Vector search branch
if op.query and self.index_config:
# We'll embed the text later, so record the request.
embedding_requests.append((idx, op.query))
score_operator, post_operator = get_distance_operator(self)
post_operator = post_operator.replace("scored", "uniq")
vector_type = (
cast(PostgresIndexConfig, self.index_config)
.get("ann_index_config", {})
.get("vector_type", "vector")
)
# For hamming bit vectors, or “regular” vectors
if (
vector_type == "bit"
and self.index_config.get("distance_type") == "hamming"
and cast(dict, self.index_config).get("distance_type") == "hamming"
):
score_operator = score_operator % (
"%s",
self.index_config["dims"],
cast(dict, self.index_config)["dims"],
)
else:
score_operator = score_operator % (
"%s",
vector_type,
)
score_operator = score_operator % ("%s", vector_type)
vectors_per_doc_estimate = self.index_config["__estimated_num_vectors"]
vectors_per_doc_estimate = cast(dict, self.index_config)[
"__estimated_num_vectors"
]
expanded_limit = (op.limit * vectors_per_doc_estimate * 2) + 1
# Vector search with CTE for proper score handling
filter_str = (
""
if not filter_conditions
else " AND " + " AND ".join(filter_conditions)
)
if op.namespace_prefix:
prefix_filter_str = f"WHERE s.prefix LIKE %s {filter_str} "
ns_args: Sequence = (f"{_namespace_to_text(op.namespace_prefix)}%",)
else:
ns_args = ()
if filter_str:
prefix_filter_str = f"WHERE {filter_str} "
else:
prefix_filter_str = ""
base_query = f"""
WITH scored AS (
SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS neg_score
FROM store s
JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
{prefix_filter_str}
ORDER BY {score_operator} ASC
# “sub_scored” does the main vector search
# Then we do DISTINCT ON to drop duplicates if your store can have them
# Finally we limit & offset
vector_search_cte = f"""
SELECT store.prefix, store.key, store.value, store.created_at, store.updated_at,
{score_operator} AS neg_score
FROM store
JOIN store_vectors sv ON store.prefix = sv.prefix AND store.key = sv.key
WHERE {ns_condition} {extra_filters}
ORDER BY {score_operator} ASC
LIMIT %s
)
SELECT * FROM (
SELECT DISTINCT ON (prefix, key)
prefix, key, value, created_at, updated_at, {post_operator} as score
FROM scored
ORDER BY prefix, key, score DESC
) AS unique_docs
ORDER BY score DESC
LIMIT %s
OFFSET %s
"""
params = [
PLACEHOLDER, # Vector placeholder
*ns_args,
"""
search_results_sql = f"""
WITH scored AS (
{vector_search_cte}
)
SELECT uniq.prefix, uniq.key, uniq.value, uniq.created_at, uniq.updated_at,
{post_operator} AS score
FROM (
SELECT DISTINCT ON (scored.prefix, scored.key)
scored.prefix, scored.key, scored.value, scored.created_at, scored.updated_at, scored.neg_score
FROM scored
ORDER BY scored.prefix, scored.key, scored.neg_score ASC
) uniq
ORDER BY score DESC
LIMIT %s
OFFSET %s
"""
search_results_params = [
PLACEHOLDER,
*ns_param,
*filter_params,
PLACEHOLDER,
expanded_limit,
@@ -440,24 +512,45 @@ class BasePostgresStore(Generic[C]):
op.offset,
]
# Regular search branch
else:
base_query = """
SELECT prefix, key, value, created_at, updated_at
FROM store
WHERE prefix LIKE %s
"""
params = [f"{_namespace_to_text(op.namespace_prefix)}%"]
base_query = f"""
SELECT store.prefix, store.key, store.value, store.created_at, store.updated_at, NULL AS score
FROM store
WHERE {ns_condition} {extra_filters}
ORDER BY store.updated_at DESC
LIMIT %s
OFFSET %s
"""
search_results_sql = base_query
search_results_params = [
*ns_param,
*filter_params,
op.limit,
op.offset,
]
if filter_conditions:
params.extend(filter_params)
base_query += " AND " + " AND ".join(filter_conditions)
base_query += " ORDER BY updated_at DESC"
base_query += " LIMIT %s OFFSET %s"
params.extend([op.limit, op.offset])
queries.append((base_query, params))
if op.refresh_ttl:
# Wrap entire primary query in a CTE, then perform "update_at"
final_sql = f"""
WITH search_results AS (
{search_results_sql}
),
updated AS (
UPDATE store s
SET expires_at = NOW() + (s.ttl_minutes || ' minutes')::interval
FROM search_results sr
WHERE s.prefix = sr.prefix
AND s.key = sr.key
AND s.ttl_minutes IS NOT NULL
)
SELECT sr.prefix, sr.key, sr.value, sr.created_at, sr.updated_at, sr.score
FROM search_results sr
"""
final_params = search_results_params[:] # copy
else:
final_sql = search_results_sql
final_params = search_results_params
queries.append((final_sql, final_params))
return queries, embedding_requests
@@ -603,6 +696,11 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
Make sure to call `setup()` before first use to create necessary tables and indexes.
The pgvector extension must be available to use vector search.
Note:
If you provide a TTL configuration, you must explicitly call `start_ttl_sweeper()` to begin
the background thread that removes expired items. Call `stop_ttl_sweeper()` to properly
clean up resources when you're done with the store.
"""
__slots__ = (
@@ -612,7 +710,10 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
"supports_pipeline",
"index_config",
"embeddings",
"_ttl_sweeper_thread",
"_ttl_stop_event",
)
supports_ttl: bool = True
def __init__(
self,
@@ -637,6 +738,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
else:
self.embeddings = None
self.ttl_config = ttl
self._ttl_sweeper_thread: Optional[threading.Thread] = None
self._ttl_stop_event = threading.Event()
@classmethod
@contextmanager
@@ -647,6 +750,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
pipeline: bool = False,
pool_config: Optional[PoolConfig] = None,
index: Optional[PostgresIndexConfig] = None,
ttl: Optional[TTLConfig] = None,
) -> Iterator["PostgresStore"]:
"""Create a new PostgresStore instance from a connection string.
@@ -678,16 +782,123 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
**cast(dict, pc),
),
) as pool:
yield cls(conn=pool, index=index)
yield cls(conn=pool, index=index, ttl=ttl)
else:
with Connection.connect(
conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
) as conn:
if pipeline:
with conn.pipeline() as pipe:
yield cls(conn, pipe=pipe, index=index)
yield cls(conn, pipe=pipe, index=index, ttl=ttl)
else:
yield cls(conn, index=index)
yield cls(conn, index=index, ttl=ttl)
def sweep_ttl(self) -> int:
"""Delete expired store items based on TTL.
Returns:
int: The number of deleted items.
"""
with self._cursor() as cur:
cur.execute(
"""
DELETE FROM store
WHERE expires_at IS NOT NULL AND expires_at < NOW()
"""
)
deleted_count = cur.rowcount
return deleted_count
def start_ttl_sweeper(
self, sweep_interval_minutes: Optional[int] = None
) -> concurrent.futures.Future[None]:
"""Periodically delete expired store items based on TTL.
Returns:
Future that can be waited on or cancelled.
"""
if not self.ttl_config:
future: concurrent.futures.Future[None] = concurrent.futures.Future()
future.set_result(None)
return future
if self._ttl_sweeper_thread and self._ttl_sweeper_thread.is_alive():
logger.info("TTL sweeper thread is already running")
# Return a future that can be used to cancel the existing thread
future = concurrent.futures.Future()
future.add_done_callback(
lambda f: self._ttl_stop_event.set() if f.cancelled() else None
)
return future
self._ttl_stop_event.clear()
interval = float(
sweep_interval_minutes or self.ttl_config.get("sweep_interval_minutes") or 5
)
logger.info(f"Starting store TTL sweeper with interval {interval} minutes")
future = concurrent.futures.Future()
def _sweep_loop() -> None:
try:
while not self._ttl_stop_event.is_set():
if self._ttl_stop_event.wait(interval * 60):
break
try:
expired_items = self.sweep_ttl()
if expired_items > 0:
logger.info(f"Store swept {expired_items} expired items")
except Exception as exc:
logger.exception(
"Store TTL sweep iteration failed", exc_info=exc
)
future.set_result(None)
except Exception as exc:
future.set_exception(exc)
thread = threading.Thread(target=_sweep_loop, daemon=True, name="ttl-sweeper")
self._ttl_sweeper_thread = thread
thread.start()
future.add_done_callback(
lambda f: self._ttl_stop_event.set() if f.cancelled() else None
)
return future
def stop_ttl_sweeper(self, timeout: Optional[float] = None) -> bool:
"""Stop the TTL sweeper thread if it's running.
Args:
timeout: Maximum time to wait for the thread to stop, in seconds.
If None, wait indefinitely.
Returns:
bool: True if the thread was successfully stopped or wasn't running,
False if the timeout was reached before the thread stopped.
"""
if not self._ttl_sweeper_thread or not self._ttl_sweeper_thread.is_alive():
return True
logger.info("Stopping TTL sweeper thread")
self._ttl_stop_event.set()
self._ttl_sweeper_thread.join(timeout)
success = not self._ttl_sweeper_thread.is_alive()
if success:
self._ttl_sweeper_thread = None
logger.info("TTL sweeper thread stopped")
else:
logger.warning("Timed out waiting for TTL sweeper thread to stop")
return success
def __del__(self) -> None:
"""Ensure the TTL sweeper thread is stopped when the object is garbage collected."""
if hasattr(self, "_ttl_stop_event") and hasattr(self, "_ttl_sweeper_thread"):
self.stop_ttl_sweeper(timeout=0.1)
@contextmanager
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
@@ -886,8 +1097,14 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
with self._cursor() as cur:
version = _get_version(cur, table="store_migrations")
for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
cur.execute(sql)
cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
try:
cur.execute(sql)
cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
except Exception as e:
logger.error(
f"Failed to apply migration {v}.\nSql={sql}\nError={e}"
)
raise
if self.index_config:
version = _get_version(cur, table="vector_migrations")
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint-postgres"
version = "2.0.16"
version = "2.0.18"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
authors = []
license = "MIT"
@@ -26,6 +26,9 @@ from tests.conftest import (
CharacterEmbeddings,
)
TTL_SECONDS = 6
TTL_MINUTES = TTL_SECONDS / 60
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
async def store(request) -> AsyncIterator[AsyncPostgresStore]:
@@ -42,28 +45,54 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
conn_string = f"{uri_base}/{database}{query_params}"
admin_conn_string = DEFAULT_URI
ttl_config = {
"default_ttl": TTL_MINUTES,
"refresh_on_read": True,
"sweep_interval_minutes": TTL_MINUTES / 2,
}
async with await AsyncConnection.connect(
admin_conn_string, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
async with AsyncPostgresStore.from_conn_string(conn_string) as store:
async with AsyncPostgresStore.from_conn_string(
conn_string, ttl=ttl_config
) as store:
store.MIGRATIONS = [
(
mig.replace("ttl_minutes INT;", "ttl_minutes FLOAT;")
if isinstance(mig, str)
else mig
)
for mig in store.MIGRATIONS
]
await store.setup()
async with store._cursor() as cur:
# drop the migration index
await cur.execute("DROP TABLE IF EXISTS store_migrations")
await store.setup() # Will fail if migrations aren't idempotent
if request.param == "pipe":
async with AsyncPostgresStore.from_conn_string(
conn_string, pipeline=True
conn_string, pipeline=True, ttl=ttl_config
) as store:
await store.start_ttl_sweeper()
yield store
await store.stop_ttl_sweeper()
elif request.param == "pool":
async with AsyncPostgresStore.from_conn_string(
conn_string, pool_config={"min_size": 1, "max_size": 10}
conn_string, pool_config={"min_size": 1, "max_size": 10}, ttl=ttl_config
) as store:
await store.start_ttl_sweeper()
yield store
await store.stop_ttl_sweeper()
else: # default
async with AsyncPostgresStore.from_conn_string(conn_string) as store:
async with AsyncPostgresStore.from_conn_string(
conn_string, ttl=ttl_config
) as store:
await store.start_ttl_sweeper()
yield store
await store.stop_ttl_sweeper()
finally:
async with await AsyncConnection.connect(
admin_conn_string, autocommit=True
@@ -635,3 +664,28 @@ async def test_search_sorting(
assert len(set(r.key for r in results)) == 10
assert results[0].key == "M"
assert results[0].score > results[1].score
async def test_store_ttl(store):
# Assumes a TTL of 1 minute = 60 seconds
ns = ("foo",)
await store.start_ttl_sweeper()
await store.aput(
ns,
key="item1",
value={"foo": "bar"},
ttl=TTL_MINUTES, # type: ignore
)
await asyncio.sleep(TTL_SECONDS - 2)
res = await store.aget(ns, key="item1", refresh_ttl=True)
assert res is not None
await asyncio.sleep(TTL_SECONDS - 2)
results = await store.asearch(ns, query="foo", refresh_ttl=True)
assert len(results) == 1
await asyncio.sleep(TTL_SECONDS - 2)
res = await store.aget(ns, key="item1", refresh_ttl=False)
assert res is not None
await asyncio.sleep(TTL_SECONDS - 1)
# Now has been (TTL_SECONDS-2)*2 > TTL_SECONDS + TTL_SECONDS/2
results = await store.asearch(ns, query="bar", refresh_ttl=False)
assert len(results) == 0
+195 -125
View File
@@ -1,6 +1,7 @@
# type: ignore
import re
import time
from contextlib import contextmanager
from typing import Any, Optional
from uuid import uuid4
@@ -24,6 +25,9 @@ from tests.conftest import (
CharacterEmbeddings,
)
TTL_SECONDS = 6
TTL_MINUTES = TTL_SECONDS / 60
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
def store(request) -> PostgresStore:
@@ -32,29 +36,56 @@ def store(request) -> PostgresStore:
uri_base = "/".join(uri_parts[:-1])
query_params = ""
if "?" in uri_parts[-1]:
db_name, query_params = uri_parts[-1].split("?", 1)
_, query_params = uri_parts[-1].split("?", 1)
query_params = "?" + query_params
conn_string = f"{uri_base}/{database}{query_params}"
admin_conn_string = DEFAULT_URI
ttl_config = {
"default_ttl": TTL_MINUTES,
"refresh_on_read": True,
"sweep_interval_minutes": TTL_MINUTES / 2,
}
with Connection.connect(admin_conn_string, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
with PostgresStore.from_conn_string(conn_string) as store:
with PostgresStore.from_conn_string(conn_string, ttl=ttl_config) as store:
store.MIGRATIONS = [
(
mig.replace("ttl_minutes INT;", "ttl_minutes FLOAT;")
if isinstance(mig, str)
else mig
)
for mig in store.MIGRATIONS
]
store.setup()
if request.param == "pipe":
with PostgresStore.from_conn_string(conn_string, pipeline=True) as store:
with PostgresStore.from_conn_string(
conn_string,
pipeline=True,
ttl=ttl_config,
) as store:
store.start_ttl_sweeper()
yield store
store.stop_ttl_sweeper()
elif request.param == "pool":
with PostgresStore.from_conn_string(
conn_string, pool_config={"min_size": 1, "max_size": 10}
conn_string,
pool_config={"min_size": 1, "max_size": 10},
ttl=ttl_config,
) as store:
store.start_ttl_sweeper()
yield store
store.stop_ttl_sweeper()
else: # default
with PostgresStore.from_conn_string(conn_string) as store:
with PostgresStore.from_conn_string(conn_string, ttl=ttl_config) as store:
store.start_ttl_sweeper()
yield store
store.stop_ttl_sweeper()
finally:
with Connection.connect(admin_conn_string, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@@ -220,134 +251,127 @@ def test_batch_list_namespaces_ops(store: PostgresStore) -> None:
assert all(ns[-1] == "public" for ns in results[2])
class TestPostgresStore:
@pytest.fixture(autouse=True)
def setup(self) -> None:
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
store.setup()
def test_basic_store_ops(store) -> None:
namespace = ("test", "documents")
item_id = "doc1"
item_value = {"title": "Test Document", "content": "Hello, World!"}
def test_basic_store_ops(self) -> None:
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
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)
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
assert item
assert item.namespace == namespace
assert item.key == item_id
assert item.value == item_value
# Test update
updated_value = {"title": "Updated Document", "content": "Hello, Updated!"}
store.put(namespace, item_id, updated_value)
updated_item = store.get(namespace, item_id)
# Test update
updated_value = {"title": "Updated Document", "content": "Hello, Updated!"}
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
assert updated_item.value == updated_value
assert updated_item.updated_at > item.updated_at
# Test get from non-existent namespace
different_namespace = ("test", "other_documents")
item_in_different_namespace = store.get(different_namespace, item_id)
assert item_in_different_namespace is None
# Test get from non-existent namespace
different_namespace = ("test", "other_documents")
item_in_different_namespace = store.get(different_namespace, item_id)
assert item_in_different_namespace is None
# Test delete
store.delete(namespace, item_id)
deleted_item = store.get(namespace, item_id)
assert deleted_item is None
# Test delete
store.delete(namespace, item_id)
deleted_item = store.get(namespace, item_id)
assert deleted_item is None
def test_list_namespaces(self) -> None:
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
# Create test data with various namespaces
test_namespaces = [
("test", "documents", "public"),
("test", "documents", "private"),
("test", "images", "public"),
("test", "images", "private"),
("prod", "documents", "public"),
("prod", "documents", "private"),
]
def test_list_namespaces(store) -> None:
# Create test data with various namespaces
test_namespaces = [
("test", "documents", "public"),
("test", "documents", "private"),
("test", "images", "public"),
("test", "images", "private"),
("prod", "documents", "public"),
("prod", "documents", "private"),
]
# Insert test data
for namespace in test_namespaces:
store.put(namespace, "dummy", {"content": "dummy"})
# Insert test data
for namespace in test_namespaces:
store.put(namespace, "dummy", {"content": "dummy"})
# Test listing with various filters
all_namespaces = store.list_namespaces()
assert len(all_namespaces) == len(test_namespaces)
# Test listing with various filters
all_namespaces = store.list_namespaces()
assert len(all_namespaces) == len(test_namespaces)
# Test prefix filtering
test_prefix_namespaces = store.list_namespaces(prefix=["test"])
assert len(test_prefix_namespaces) == 4
assert all(ns[0] == "test" for ns in test_prefix_namespaces)
# Test prefix filtering
test_prefix_namespaces = store.list_namespaces(prefix=["test"])
assert len(test_prefix_namespaces) == 4
assert all(ns[0] == "test" for ns in test_prefix_namespaces)
# Test suffix filtering
public_namespaces = store.list_namespaces(suffix=["public"])
assert len(public_namespaces) == 3
assert all(ns[-1] == "public" for ns in public_namespaces)
# Test suffix filtering
public_namespaces = store.list_namespaces(suffix=["public"])
assert len(public_namespaces) == 3
assert all(ns[-1] == "public" for ns in public_namespaces)
# Test max depth
depth_2_namespaces = store.list_namespaces(max_depth=2)
assert all(len(ns) <= 2 for ns in depth_2_namespaces)
# Test max depth
depth_2_namespaces = store.list_namespaces(max_depth=2)
assert all(len(ns) <= 2 for ns in depth_2_namespaces)
# Test pagination
paginated_namespaces = store.list_namespaces(limit=3)
assert len(paginated_namespaces) == 3
# Test pagination
paginated_namespaces = store.list_namespaces(limit=3)
assert len(paginated_namespaces) == 3
# Cleanup
for namespace in test_namespaces:
store.delete(namespace, "dummy")
# Cleanup
for namespace in test_namespaces:
store.delete(namespace, "dummy")
def test_search(self) -> None:
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
# Create test data
test_data = [
(
("test", "docs"),
"doc1",
{"title": "First Doc", "author": "Alice", "tags": ["important"]},
),
(
("test", "docs"),
"doc2",
{"title": "Second Doc", "author": "Bob", "tags": ["draft"]},
),
(
("test", "images"),
"img1",
{"title": "Image 1", "author": "Alice", "tags": ["final"]},
),
]
for namespace, key, value in test_data:
store.put(namespace, key, value)
def test_search(store) -> None:
# Create test data
test_data = [
(
("test", "docs"),
"doc1",
{"title": "First Doc", "author": "Alice", "tags": ["important"]},
),
(
("test", "docs"),
"doc2",
{"title": "Second Doc", "author": "Bob", "tags": ["draft"]},
),
(
("test", "images"),
"img1",
{"title": "Image 1", "author": "Alice", "tags": ["final"]},
),
]
# Test basic search
all_items = store.search(["test"])
assert len(all_items) == 3
for namespace, key, value in test_data:
store.put(namespace, key, value)
# Test namespace filtering
docs_items = store.search(["test", "docs"])
assert len(docs_items) == 2
assert all(item.namespace == ("test", "docs") for item in docs_items)
# Test basic search
all_items = store.search(["test"])
assert len(all_items) == 3
# Test value filtering
alice_items = store.search(["test"], filter={"author": "Alice"})
assert len(alice_items) == 2
assert all(item.value["author"] == "Alice" for item in alice_items)
# Test namespace filtering
docs_items = store.search(["test", "docs"])
assert len(docs_items) == 2
assert all(item.namespace == ("test", "docs") for item in docs_items)
# Test pagination
paginated_items = store.search(["test"], limit=2)
assert len(paginated_items) == 2
# Test value filtering
alice_items = store.search(["test"], filter={"author": "Alice"})
assert len(alice_items) == 2
assert all(item.value["author"] == "Alice" for item in alice_items)
offset_items = store.search(["test"], offset=2)
assert len(offset_items) == 1
# Test pagination
paginated_items = store.search(["test"], limit=2)
assert len(paginated_items) == 2
# Cleanup
for namespace, key, _ in test_data:
store.delete(namespace, key)
offset_items = store.search(["test"], offset=2)
assert len(offset_items) == 1
# Cleanup
for namespace, key, _ in test_data:
store.delete(namespace, key)
@contextmanager
@@ -356,6 +380,7 @@ def _create_vector_store(
distance_type: str,
fake_embeddings: Embeddings,
text_fields: Optional[list[str]] = None,
enable_ttl: bool = True,
) -> PostgresStore:
"""Create a store with vector search enabled."""
database = f"test_{uuid4().hex[:16]}"
@@ -385,23 +410,32 @@ def _create_vector_store(
with PostgresStore.from_conn_string(
conn_string,
index=index_config,
ttl={"default_ttl": 2, "refresh_on_read": True} if enable_ttl else None,
) as store:
store.setup()
with store._cursor() as cur:
# drop the migration index
cur.execute("DROP TABLE IF EXISTS store_migrations")
store.setup() # Will fail if migrations aren't idempotent
yield store
finally:
with Connection.connect(admin_conn_string, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
_vector_params = [
(vector_type, distance_type, True)
for vector_type in VECTOR_TYPES
for distance_type in (
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
)
]
_vector_params += [(*_vector_params[-1][:2], False)]
@pytest.fixture(
scope="function",
params=[
(vector_type, distance_type)
for vector_type in VECTOR_TYPES
for distance_type in (
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
)
],
params=_vector_params,
ids=lambda p: f"{p[0]}_{p[1]}",
)
def vector_store(
@@ -409,8 +443,10 @@ def vector_store(
fake_embeddings: Embeddings,
) -> PostgresStore:
"""Create a store with vector search enabled."""
vector_type, distance_type = request.param
with _create_vector_store(vector_type, distance_type, fake_embeddings) as store:
vector_type, distance_type, enable_ttl = request.param
with _create_vector_store(
vector_type, distance_type, fake_embeddings, enable_ttl=enable_ttl
) as store:
yield store
@@ -474,7 +510,10 @@ def test_vector_update_with_embedding(vector_store: PostgresStore) -> None:
assert not any(r.key == "doc4" for r in results_new)
def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
@pytest.mark.parametrize("refresh_ttl", [True, False])
def test_vector_search_with_filters(
vector_store: PostgresStore, refresh_ttl: bool
) -> None:
"""Test combining vector search with filters."""
# Insert test documents
docs = [
@@ -487,16 +526,23 @@ def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
for key, value in docs:
vector_store.put(("test",), key, value)
results = vector_store.search(("test",), query="apple", filter={"color": "red"})
results = vector_store.search(
("test",), query="apple", filter={"color": "red"}, refresh_ttl=refresh_ttl
)
assert len(results) == 2
assert results[0].key == "doc1"
results = vector_store.search(("test",), query="car", filter={"color": "red"})
results = vector_store.search(
("test",), query="car", filter={"color": "red"}, refresh_ttl=refresh_ttl
)
assert len(results) == 2
assert results[0].key == "doc2"
results = vector_store.search(
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
("test",),
query="bbbbluuu",
filter={"score": {"$gt": 3.2}},
refresh_ttl=refresh_ttl,
)
assert len(results) == 3
assert results[0].key == "doc4"
@@ -688,7 +734,7 @@ def test_embed_with_path_operation_config(
store.put(("test",), "doc5", doc5, index=False)
results = store.search(("test",))
assert len(results) == 3
assert all(r.score is None for r in results)
assert all(r.score is None for r in results), f"{results}"
assert any(r.key == "doc5" for r in results)
results = store.search(("test",), query="hhh")
@@ -790,3 +836,27 @@ def test_nonnull_migrations() -> None:
for migration in PostgresStore.MIGRATIONS:
statement = _leading_comment_remover.sub("", migration).split()[0]
assert statement.strip()
def test_store_ttl(store):
# Assumes a TTL of 1 minute = 60 seconds
ns = ("foo",)
store.put(
ns,
key="item1",
value={"foo": "bar"},
ttl=TTL_MINUTES, # type: ignore
)
time.sleep(TTL_SECONDS - 2)
res = store.get(ns, key="item1", refresh_ttl=True)
assert res is not None
time.sleep(TTL_SECONDS - 2)
results = store.search(ns, query="foo", refresh_ttl=True)
assert len(results) == 1
time.sleep(TTL_SECONDS - 2)
res = store.get(ns, key="item1", refresh_ttl=False)
assert res is not None
time.sleep(TTL_SECONDS - 1)
# Now has been (TTL_SECONDS-2)*2 > TTL_SECONDS + TTL_SECONDS/2
res = store.search(ns, query="bar", refresh_ttl=False)
assert len(res) == 0
@@ -45,3 +45,18 @@ def maybe_add_typed_methods(serde: SerializerProtocol) -> SerializerProtocol:
return SerializerCompat(serde)
return serde
class CipherProtocol(Protocol):
"""Protocol for encryption and decryption of data.
- `encrypt`: Encrypt plaintext.
- `decrypt`: Decrypt ciphertext.
"""
def encrypt(self, plaintext: bytes) -> tuple[str, bytes]:
"""Encrypt plaintext. Returns a tuple (cipher name, ciphertext)."""
...
def decrypt(self, ciphername: str, ciphertext: bytes) -> bytes:
"""Decrypt ciphertext. Returns the plaintext."""
...
@@ -0,0 +1,86 @@
import os
from typing import Any
from langgraph.checkpoint.serde.base import CipherProtocol, SerializerProtocol
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
class EncryptedSerializer(SerializerProtocol):
"""Serializer that encrypts and decrypts data using an encryption protocol."""
def __init__(
self, cipher: CipherProtocol, serde: SerializerProtocol = JsonPlusSerializer()
) -> None:
self.cipher = cipher
self.serde = serde
def dumps(self, obj: Any) -> bytes:
return self.serde.dumps(obj)
def loads(self, data: bytes) -> Any:
return self.serde.loads(data)
def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
"""Serialize an object to a tuple (type, bytes) and encrypt the bytes."""
# serialize data
typ, data = self.serde.dumps_typed(obj)
# encrypt data
ciphername, ciphertext = self.cipher.encrypt(data)
# add cipher name to type
return f"{typ}+{ciphername}", ciphertext
def loads_typed(self, data: tuple[str, bytes]) -> Any:
enc_cipher, ciphertext = data
# unencrypted data
if "+" not in enc_cipher:
return self.serde.loads_typed(data)
# extract cipher name
typ, ciphername = enc_cipher.split("+", 1)
# decrypt data
decrypted_data = self.cipher.decrypt(ciphername, ciphertext)
# deserialize data
return self.serde.loads_typed((typ, decrypted_data))
@classmethod
def from_pycryptodome_aes(
cls, serde: SerializerProtocol = JsonPlusSerializer(), **kwargs: Any
) -> "EncryptedSerializer":
"""Create an EncryptedSerializer using AES encryption."""
try:
from Crypto.Cipher import AES # type: ignore
except ImportError:
raise ImportError(
"Pycryptodome is not installed. Please install it with `pip install pycryptodome`."
) from None
# check if AES key is provided
if "key" in kwargs:
key: bytes = kwargs.pop("key")
else:
key_str = os.getenv("LANGGRAPH_AES_KEY")
if key_str is None:
raise ValueError("LANGGRAPH_AES_KEY environment variable is not set.")
key = key_str.encode()
if len(key) not in (16, 24, 32):
raise ValueError("LANGGRAPH_AES_KEY must be 16, 24, or 32 bytes long.")
# set default mode to EAX if not provided
if kwargs.get("mode") is None:
kwargs["mode"] = AES.MODE_EAX
class PycryptodomeAesCipher(CipherProtocol):
def encrypt(self, plaintext: bytes) -> tuple[str, bytes]:
cipher = AES.new(key, **kwargs)
ciphertext, tag = cipher.encrypt_and_digest(plaintext)
return "aes", cipher.nonce + tag + ciphertext
def decrypt(self, ciphername: str, ciphertext: bytes) -> bytes:
assert ciphername == "aes", f"Unsupported cipher: {ciphername}"
nonce = ciphertext[:16]
tag = ciphertext[16:32]
actual_ciphertext = ciphertext[32:]
cipher = AES.new(key, **kwargs, nonce=nonce)
return cipher.decrypt_and_verify(actual_ciphertext, tag)
return cls(PycryptodomeAesCipher(), serde)
@@ -537,6 +537,12 @@ class TTLConfig(TypedDict, total=False):
The expiration timer refreshes on both read and write operations.
Defaults to None (no expiration).
"""
sweep_interval_minutes: Optional[int]
"""Interval in minutes between TTL sweep operations.
If provided, the store will periodically delete expired items based on TTL.
Defaults to None (no sweeping).
"""
class IndexConfig(TypedDict, total=False):
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint"
version = "2.0.19"
version = "2.0.21"
description = "Library with base interfaces for LangGraph checkpoint savers."
authors = []
license = "MIT"
+4 -1
View File
@@ -1,4 +1,4 @@
.PHONY: test lint format test-integration
.PHONY: test lint format test-integration update-schema
######################
# TESTING AND COVERAGE
@@ -31,3 +31,6 @@ lint lint_diff lint_package lint_tests:
format format_diff:
poetry run ruff format $(PYTHON_FILES)
poetry run ruff check --select I --fix $(PYTHON_FILES)
update-schema:
poetry run python generate_schema.py
+6
View File
@@ -27,6 +27,12 @@ class TTLConfig(TypedDict, total=False):
If provided, all new items will have this TTL unless explicitly overridden.
If omitted, items will have no TTL by default.
"""
sweep_interval_minutes: Optional[int]
"""Optional. Interval in minutes between TTL sweep iterations.
If provided, the store will periodically delete expired items based on the TTL.
If omitted, no automatic sweeping will occur.
"""
class IndexConfig(TypedDict, total=False):
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-cli"
version = "0.1.76"
version = "0.1.77"
description = "CLI for interacting with LangGraph API"
authors = []
license = "MIT"
+10
View File
@@ -459,6 +459,16 @@
},
"refresh_on_read": {
"type": "boolean"
},
"sweep_interval_minutes": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
]
}
},
"required": []
+10
View File
@@ -459,6 +459,16 @@
},
"refresh_on_read": {
"type": "boolean"
},
"sweep_interval_minutes": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
]
}
},
"required": []
+1 -1
View File
@@ -416,7 +416,7 @@ class StateGraph(Graph):
and (vals := get_args(rargs[0]))
):
ends = vals
except (TypeError, StopIteration):
except (NameError, TypeError, StopIteration):
pass
if destinations is not None:
+44 -2
View File
@@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 2.0.1 and should not be changed by hand.
# This file is automatically @generated by Poetry 2.0.0 and should not be changed by hand.
[[package]]
name = "aiosqlite"
@@ -2238,6 +2238,48 @@ files = [
{file = "pycparser-2.22.tar.gz", hash = "sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6"},
]
[[package]]
name = "pycryptodome"
version = "3.21.0"
description = "Cryptographic library for Python"
optional = false
python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7"
groups = ["dev"]
files = [
{file = "pycryptodome-3.21.0-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:dad9bf36eda068e89059d1f07408e397856be9511d7113ea4b586642a429a4fd"},
{file = "pycryptodome-3.21.0-cp27-cp27m-manylinux2010_i686.whl", hash = "sha256:a1752eca64c60852f38bb29e2c86fca30d7672c024128ef5d70cc15868fa10f4"},
{file = "pycryptodome-3.21.0-cp27-cp27m-manylinux2010_x86_64.whl", hash = "sha256:3ba4cc304eac4d4d458f508d4955a88ba25026890e8abff9b60404f76a62c55e"},
{file = "pycryptodome-3.21.0-cp27-cp27m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7cb087b8612c8a1a14cf37dd754685be9a8d9869bed2ffaaceb04850a8aeef7e"},
{file = "pycryptodome-3.21.0-cp27-cp27m-musllinux_1_1_aarch64.whl", hash = "sha256:26412b21df30b2861424a6c6d5b1d8ca8107612a4cfa4d0183e71c5d200fb34a"},
{file = "pycryptodome-3.21.0-cp27-cp27m-win32.whl", hash = "sha256:cc2269ab4bce40b027b49663d61d816903a4bd90ad88cb99ed561aadb3888dd3"},
{file = "pycryptodome-3.21.0-cp27-cp27m-win_amd64.whl", hash = "sha256:0fa0a05a6a697ccbf2a12cec3d6d2650b50881899b845fac6e87416f8cb7e87d"},
{file = "pycryptodome-3.21.0-cp27-cp27mu-manylinux2010_i686.whl", hash = "sha256:6cce52e196a5f1d6797ff7946cdff2038d3b5f0aba4a43cb6bf46b575fd1b5bb"},
{file = "pycryptodome-3.21.0-cp27-cp27mu-manylinux2010_x86_64.whl", hash = "sha256:a915597ffccabe902e7090e199a7bf7a381c5506a747d5e9d27ba55197a2c568"},
{file = "pycryptodome-3.21.0-cp27-cp27mu-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a4e74c522d630766b03a836c15bff77cb657c5fdf098abf8b1ada2aebc7d0819"},
{file = "pycryptodome-3.21.0-cp27-cp27mu-musllinux_1_1_aarch64.whl", hash = "sha256:a3804675283f4764a02db05f5191eb8fec2bb6ca34d466167fc78a5f05bbe6b3"},
{file = "pycryptodome-3.21.0-cp36-abi3-macosx_10_9_universal2.whl", hash = "sha256:2480ec2c72438430da9f601ebc12c518c093c13111a5c1644c82cdfc2e50b1e4"},
{file = "pycryptodome-3.21.0-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:de18954104667f565e2fbb4783b56667f30fb49c4d79b346f52a29cb198d5b6b"},
{file = "pycryptodome-3.21.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2de4b7263a33947ff440412339cb72b28a5a4c769b5c1ca19e33dd6cd1dcec6e"},
{file = "pycryptodome-3.21.0-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0714206d467fc911042d01ea3a1847c847bc10884cf674c82e12915cfe1649f8"},
{file = "pycryptodome-3.21.0-cp36-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7d85c1b613121ed3dbaa5a97369b3b757909531a959d229406a75b912dd51dd1"},
{file = "pycryptodome-3.21.0-cp36-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:8898a66425a57bcf15e25fc19c12490b87bd939800f39a03ea2de2aea5e3611a"},
{file = "pycryptodome-3.21.0-cp36-abi3-musllinux_1_2_i686.whl", hash = "sha256:932c905b71a56474bff8a9c014030bc3c882cee696b448af920399f730a650c2"},
{file = "pycryptodome-3.21.0-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:18caa8cfbc676eaaf28613637a89980ad2fd96e00c564135bf90bc3f0b34dd93"},
{file = "pycryptodome-3.21.0-cp36-abi3-win32.whl", hash = "sha256:280b67d20e33bb63171d55b1067f61fbd932e0b1ad976b3a184303a3dad22764"},
{file = "pycryptodome-3.21.0-cp36-abi3-win_amd64.whl", hash = "sha256:b7aa25fc0baa5b1d95b7633af4f5f1838467f1815442b22487426f94e0d66c53"},
{file = "pycryptodome-3.21.0-pp27-pypy_73-manylinux2010_x86_64.whl", hash = "sha256:2cb635b67011bc147c257e61ce864879ffe6d03342dc74b6045059dfbdedafca"},
{file = "pycryptodome-3.21.0-pp27-pypy_73-win32.whl", hash = "sha256:4c26a2f0dc15f81ea3afa3b0c87b87e501f235d332b7f27e2225ecb80c0b1cdd"},
{file = "pycryptodome-3.21.0-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:d5ebe0763c982f069d3877832254f64974139f4f9655058452603ff559c482e8"},
{file = "pycryptodome-3.21.0-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7ee86cbde706be13f2dec5a42b52b1c1d1cbb90c8e405c68d0755134735c8dc6"},
{file = "pycryptodome-3.21.0-pp310-pypy310_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0fd54003ec3ce4e0f16c484a10bc5d8b9bd77fa662a12b85779a2d2d85d67ee0"},
{file = "pycryptodome-3.21.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:5dfafca172933506773482b0e18f0cd766fd3920bd03ec85a283df90d8a17bc6"},
{file = "pycryptodome-3.21.0-pp39-pypy39_pp73-macosx_10_15_x86_64.whl", hash = "sha256:590ef0898a4b0a15485b05210b4a1c9de8806d3ad3d47f74ab1dc07c67a6827f"},
{file = "pycryptodome-3.21.0-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f35e442630bc4bc2e1878482d6f59ea22e280d7121d7adeaedba58c23ab6386b"},
{file = "pycryptodome-3.21.0-pp39-pypy39_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ff99f952db3db2fbe98a0b355175f93ec334ba3d01bbde25ad3a5a33abc02b58"},
{file = "pycryptodome-3.21.0-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:8acd7d34af70ee63f9a849f957558e49a98f8f1634f86a59d2be62bb8e93f71c"},
{file = "pycryptodome-3.21.0.tar.gz", hash = "sha256:f7787e0d469bdae763b876174cf2e6c0f7be79808af26b1da96f1a64bcf47297"},
]
[[package]]
name = "pydantic"
version = "2.9.2"
@@ -3509,4 +3551,4 @@ type = ["pytest-mypy"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.9.0,<4.0"
content-hash = "eb85f0bcc0e8a715ef38afb58cf888f7c2ee8579ea6ed94900244365f24cddd9"
content-hash = "b8641a0b2d92bee0363602e69f99b23366b2035b7e17ff017708194e6fbd0ac5"
+2 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph"
version = "0.3.10"
version = "0.3.12"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
license = "MIT"
@@ -37,6 +37,7 @@ uvloop = "0.21.0beta1"
pyperf = "^2.7.0"
py-spy = "^0.3.14"
types-requests = "^2.32.0.20240914"
pycryptodome = "^3.21.0"
[tool.ruff]
lint.select = [ "E", "F", "I", "TID251" ]
File diff suppressed because one or more lines are too long
@@ -377,6 +377,19 @@
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge[sqlite_aes]
'''
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_pydantic1[memory]
'''
graph TD;
@@ -797,6 +810,76 @@
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite_aes]
'''
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_pydantic1[sqlite_aes].1
dict({
'definitions': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'inner': dict({
'$ref': '#/definitions/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
]),
'title': 'Input',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite_aes].2
dict({
'properties': dict({
'answer': dict({
'title': 'Answer',
'type': 'string',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory]
'''
graph TD;
@@ -1217,6 +1300,76 @@
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aes]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aes].1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
]),
'title': 'Input',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aes].2
dict({
'properties': dict({
'answer': dict({
'title': 'Answer',
'type': 'string',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[memory]
'''
graph TD;
@@ -1715,6 +1868,19 @@
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[sqlite_aes]
'''
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_multiple_sinks_subgraphs
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
+11
View File
@@ -16,6 +16,7 @@ from langgraph.checkpoint.postgres.aio import (
AsyncPostgresSaver,
AsyncShallowPostgresSaver,
)
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.store.base import BaseStore
@@ -61,6 +62,15 @@ def checkpointer_sqlite():
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_sqlite_aes():
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes(
key=b"1234567890123456"
)
yield checkpointer
@asynccontextmanager
async def _checkpointer_sqlite_aio():
async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
@@ -437,6 +447,7 @@ REGULAR_CHECKPOINTERS_SYNC = [
"postgres",
"postgres_pipe",
"postgres_pool",
"sqlite_aes",
]
ALL_CHECKPOINTERS_SYNC = [
*REGULAR_CHECKPOINTERS_SYNC,