mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-29 03:09:45 +02:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2e1adaa867 | ||
|
|
3f8b165592 | ||
|
|
9ed0fa196c | ||
|
|
0b9adc28c3 | ||
|
|
3b0255d1ef | ||
|
|
d4255a0645 | ||
|
|
80d61a2600 | ||
|
|
424f24720a | ||
|
|
2a71180c1d | ||
|
|
697f878e36 | ||
|
|
def69c59d2 | ||
|
|
bad4d17c34 | ||
|
|
55219b23d8 | ||
|
|
8edbd39ad3 | ||
|
|
4b0fd834d8 | ||
|
|
0fd2748530 | ||
|
|
bc0a3419ed | ||
|
|
5cd47bac49 | ||
|
|
4c6d80a67f | ||
|
|
18ed044c27 | ||
|
|
394a9fa85f | ||
|
|
9741d9bdf0 | ||
|
|
06ca07432d |
@@ -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
|
||||
@@ -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,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
|
||||
|
||||
@@ -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,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
@@ -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
|
||||
|
||||
@@ -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,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"
|
||||
|
||||
@@ -459,6 +459,16 @@
|
||||
},
|
||||
"refresh_on_read": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"sweep_interval_minutes": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"required": []
|
||||
|
||||
@@ -459,6 +459,16 @@
|
||||
},
|
||||
"refresh_on_read": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"sweep_interval_minutes": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"required": []
|
||||
|
||||
@@ -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:
|
||||
|
||||
Generated
+44
-2
@@ -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"
|
||||
|
||||
@@ -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'}}}%%
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user