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

Expose filter arg in get_state_history

Browse Source 
- Combine list and search methods in Checkpointer
This commit is contained in:
Nuno Campos
2024-06-05 17:14:35 -07:00
parent 052284f21d
commit d366bacf03
11 changed files with 154 additions and 348 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/checkpoint/aiosqlite.py
+11 -85
View File
@@ -2,7 +2,7 @@ import asyncio
import functools
from contextlib import AbstractAsyncContextManager
from types import TracebackType
from typing import AsyncIterator, Iterator, Optional, TypeVar
from typing import Any, AsyncIterator, Dict, Iterator, Optional, TypeVar
import aiosqlite
from langchain_core.runnables import RunnableConfig
@@ -141,8 +141,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
@not_implemented_sync_method
def list(
self,
config: RunnableConfig,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
@@ -153,21 +154,6 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
Or consider using the [SqliteSaver](#sqlitesaver) checkpointer.
"""
@not_implemented_sync_method
def search(
self,
metadata_filter: CheckpointMetadata,
*,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
"""Search for checkpoints by metadata.
Note:
This method is not implemented for the AsyncSqliteSaver. Use `asearch` instead.
Or consider using the [SqliteSaver](#sqlitesaver) checkpointer.
"""
@not_implemented_sync_method
def put(
self,
@@ -274,8 +260,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
async def alist(
self,
config: RunnableConfig,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
@@ -293,75 +280,14 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples.
"""
await self.setup()
query = (
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC"
if before is None
else "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts < ? ORDER BY thread_ts DESC"
)
where, param_values = search_where(config, filter, before)
query = f"""SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata
FROM checkpoints
{where}
ORDER BY thread_ts DESC"""
if limit:
query += f" LIMIT {limit}"
async with self.conn.execute(
query,
(
(str(config["configurable"]["thread_id"]),)
if before is None
else (
str(config["configurable"]["thread_id"]),
str(before["configurable"]["thread_ts"]),
)
),
) as cursor:
async for thread_id, thread_ts, parent_ts, value, metadata in cursor:
yield CheckpointTuple(
{"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
self.serde.loads(value),
self.serde.loads(metadata) if metadata is not None else {},
(
{
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
}
if parent_ts
else None
),
)
async def asearch(
self,
metadata_filter: CheckpointMetadata,
*,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
"""Search for checkpoints by metadata asynchronously.
This method retrieves a list of checkpoint tuples from the SQLite
database based on the provided metadata filter. The metadata filter does
not need to contain all keys defined in the CheckpointMetadata class.
The checkpoints are ordered by timestamp in descending order.
Args:
metadata_filter (CheckpointMetadata): The metadata filter to use for searching the checkpoints.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None.
limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
Yields:
Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
"""
await self.setup()
# construct query
SELECT = "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints "
WHERE, params = search_where(metadata_filter, before)
ORDER_BY = "ORDER BY thread_ts DESC "
LIMIT = f"LIMIT {limit}" if limit else ""
query = f"{SELECT}{WHERE}{ORDER_BY}{LIMIT}"
# execute query
async with self.conn.execute(query, params) as cursor:
async with self.conn.execute(query, param_values) as cursor:
async for thread_id, thread_ts, parent_ts, value, metadata in cursor:
yield CheckpointTuple(
{"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
langgraph/checkpoint/base.py
+4 -21
View File
@@ -4,6 +4,7 @@ from datetime import datetime, timezone
from typing import (
Any,
AsyncIterator,
Dict,
Iterator,
List,
Literal,
@@ -159,17 +160,9 @@ class BaseCheckpointSaver(ABC):
def list(
self,
config: RunnableConfig,
*,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
raise NotImplementedError
def search(
self,
metadata_filter: CheckpointMetadata,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
@@ -192,17 +185,7 @@ class BaseCheckpointSaver(ABC):
def alist(
self,
config: RunnableConfig,
*,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
raise NotImplementedError
yield
def asearch(
self,
metadata_filter: CheckpointMetadata,
config: Optional[RunnableConfig],
*,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
langgraph/checkpoint/memory.py
+26 -85
View File
@@ -1,7 +1,7 @@
import asyncio
from collections import defaultdict
from functools import partial
from typing import AsyncIterator, Iterator, Optional
from typing import Any, AsyncIterator, Dict, Iterator, Optional
from langchain_core.runnables import RunnableConfig
@@ -89,8 +89,9 @@ class MemorySaver(BaseCheckpointSaver):
def list(
self,
config: RunnableConfig,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
@@ -107,71 +108,33 @@ class MemorySaver(BaseCheckpointSaver):
Yields:
Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
"""
thread_id = config["configurable"]["thread_id"]
for ts, (checkpoint, metadata) in self.storage[thread_id].items():
if before and ts >= before["configurable"]["thread_ts"]:
continue
if limit is not None and limit <= 0:
break
elif limit is not None:
limit -= 1
yield CheckpointTuple(
config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
checkpoint=self.serde.loads(checkpoint),
metadata=self.serde.loads(metadata),
)
def search(
self,
metadata_filter: CheckpointMetadata,
*,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
"""Search for checkpoints by metadata.
This method retrieves a list of checkpoint tuples from the in-memory
storage based on the provided metadata filter. The metadata filter does
not need to contain all keys defined in the CheckpointMetadata class.
The checkpoints are ordered by timestamp in descending order.
Args:
metadata_filter (CheckpointMetadata): The metadata filter to use for searching the checkpoints.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None.
limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
Yields:
Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
"""
for thread_id, checkpoints in self.storage.items():
for ts, (checkpoint_bytes, metadata_bytes) in checkpoints.items():
thread_ids = (config["configurable"]["thread_id"],) if config else self.storage
for thread_id in thread_ids:
for ts, (checkpoint, metadata_b) in self.storage[thread_id].items():
# filter by thread_ts
if before and ts >= before["configurable"]["thread_ts"]:
continue
# check if all query key/value pairs match the metadata
metadata = self.serde.loads(metadata_bytes)
all_keys_match = all(
# filter by metadata
metadata = self.serde.loads(metadata_b)
if filter and not all(
query_value == metadata[query_key]
for query_key, query_value in metadata_filter.items()
for query_key, query_value in filter.items()
):
continue
# limit search results
if limit is not None and limit <= 0:
break
elif limit is not None:
limit -= 1
yield CheckpointTuple(
config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
checkpoint=self.serde.loads(checkpoint),
metadata=metadata,
)
# if all query key/value pairs match, yield the checkpoint
if all_keys_match:
# limit search results
if limit is not None:
if limit <= 0:
break
limit -= 1
yield CheckpointTuple(
config={
"configurable": {"thread_id": thread_id, "thread_ts": ts}
},
checkpoint=self.serde.loads(checkpoint_bytes),
metadata=metadata,
)
def put(
self,
config: RunnableConfig,
@@ -223,8 +186,9 @@ class MemorySaver(BaseCheckpointSaver):
async def alist(
self,
config: RunnableConfig,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
@@ -241,7 +205,7 @@ class MemorySaver(BaseCheckpointSaver):
"""
loop = asyncio.get_running_loop()
iter = await loop.run_in_executor(
None, partial(self.list, before=before, limit=limit), config
None, partial(self.list, before=before, limit=limit, filter=filter), config
)
while True:
# handling StopIteration exception inside coroutine won't work
@@ -251,29 +215,6 @@ class MemorySaver(BaseCheckpointSaver):
else:
break
async def asearch(
self,
metadata_filter: CheckpointMetadata,
*,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
"""Asynchronous version of search.
This method is an asynchronous wrapper around search that runs the synchronous
method in a separate thread using asyncio.
"""
loop = asyncio.get_running_loop()
iter = await loop.run_in_executor(
None, partial(self.search, before=before, limit=limit), metadata_filter
)
while True:
if item := await loop.run_in_executor(None, next, iter, None):
yield item
else:
break
async def aput(
self,
config: RunnableConfig,
langgraph/checkpoint/sqlite.py
+36 -118
View File
@@ -4,7 +4,7 @@ import sqlite3
import threading
from contextlib import AbstractContextManager, contextmanager
from types import TracebackType
from typing import Any, AsyncIterator, Iterator, Optional, Tuple
from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
@@ -283,8 +283,9 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
def list(
self,
config: RunnableConfig,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
@@ -316,76 +317,15 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
>>> print(checkpoints)
[CheckpointTuple(...), ...]
"""
query = (
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC"
if before is None
else "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts < ? ORDER BY thread_ts DESC"
)
where, param_values = search_where(config, filter, before)
query = f"""SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata
FROM checkpoints
{where}
ORDER BY thread_ts DESC"""
if limit:
query += f" LIMIT {limit}"
with self.cursor(transaction=False) as cur:
cur.execute(
query,
(
(str(config["configurable"]["thread_id"]),)
if before is None
else (
str(config["configurable"]["thread_id"]),
before["configurable"]["thread_ts"],
)
),
)
for thread_id, thread_ts, parent_ts, value, metadata in cur:
yield CheckpointTuple(
{"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
self.serde.loads(value),
self.serde.loads(metadata) if metadata is not None else {},
(
{
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
}
if parent_ts
else None
),
)
def search(
self,
metadata_filter: CheckpointMetadata,
*,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
"""Search for checkpoints by metadata.
This method retrieves a list of checkpoint tuples from the SQLite
database based on the provided metadata filter. The metadata filter does
not need to contain all keys defined in the CheckpointMetadata class.
The checkpoints are ordered by timestamp in descending order.
Args:
metadata_filter (CheckpointMetadata): The metadata filter to use for searching the checkpoints.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None.
limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
Yields:
Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
"""
# construct query
SELECT = "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints "
WHERE, params = search_where(metadata_filter, before)
ORDER_BY = "ORDER BY thread_ts DESC "
LIMIT = f"LIMIT {limit}" if limit else ""
query = f"{SELECT}{WHERE}{ORDER_BY}{LIMIT}"
# execute query
with self.cursor(transaction=False) as cur:
cur.execute(query, params)
cur.execute(query, param_values)
for thread_id, thread_ts, parent_ts, value, metadata in cur:
yield CheckpointTuple(
{"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
@@ -462,8 +402,9 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
async def alist(
self,
config: RunnableConfig,
config: Optional[RunnableConfig],
*,
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
@@ -476,22 +417,6 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
raise NotImplementedError(_AIO_ERROR_MSG)
yield
async def asearch(
self,
metadata_filter: CheckpointMetadata,
*,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
"""Search for checkpoints by metadata asynchronously.
Note:
This async method is not supported by the SqliteSaver class.
Use search() instead, or consider using [AsyncSqliteSaver](#asyncsqlitesaver).
"""
raise NotImplementedError(_AIO_ERROR_MSG)
yield
async def aput(
self,
config: RunnableConfig,
@@ -508,8 +433,8 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
def _metadata_predicate(
metadata_filter: CheckpointMetadata,
) -> Tuple[str, Tuple[Any, ...]]:
metadata_filter: Dict[str, Any],
) -> Tuple[Sequence[str], Sequence[Any]]:
"""Return WHERE clause predicates for (a)search() given metadata filter.
This method returns a tuple of a string and a tuple of values. The string
@@ -537,29 +462,25 @@ def _metadata_predicate(
else:
return ("= ?", str(query_value))
predicate = ""
param_values = ()
predicates = []
param_values = []
# process metadata query
for query_key, query_value in metadata_filter.items():
operator, param_value = _where_value(query_value)
predicate += (
f"json_extract(CAST(metadata AS TEXT), '$.{query_key}') {operator} AND "
predicates.append(
f"json_extract(CAST(metadata AS TEXT), '$.{query_key}') {operator}"
)
param_values += (param_value,)
param_values.append(param_value)
if predicate != "":
# remove trailing AND
predicate = predicate[:-4]
# predicate contains an extra trailing space
return (predicate, param_values)
return (predicates, param_values)
def search_where(
metadata_filter: CheckpointMetadata,
config: Optional[RunnableConfig],
filter: Optional[Dict[str, Any]],
before: Optional[RunnableConfig] = None,
) -> Tuple[str, Tuple[Any, ...]]:
) -> Tuple[str, Sequence[Any]]:
"""Return WHERE clause predicates for (a)search() given metadata filter
and `before` config.
@@ -568,26 +489,23 @@ def search_where(
"WHERE column1 = ? AND column2 IS ?". The tuple of values contains the
values for each of the corresponding parameters.
"""
where = "WHERE "
param_values = ()
wheres = []
param_values = []
# construct predicate for config filter
if config is not None:
wheres.append("thread_id = ?")
param_values.append(config["configurable"]["thread_id"])
# construct predicate for metadata filter
metadata_predicate, metadata_values = _metadata_predicate(metadata_filter)
if metadata_predicate != "":
where += metadata_predicate
param_values += metadata_values
if filter:
metadata_predicates, metadata_values = _metadata_predicate(filter)
wheres.extend(metadata_predicates)
param_values.extend(metadata_values)
# construct predicate for `before`
if before is not None:
if metadata_predicate != "":
where += "AND thread_ts < ? "
else:
where += "thread_ts < ? "
wheres.append("thread_ts < ?")
param_values.append(before["configurable"]["thread_ts"])
param_values += (before["configurable"]["thread_ts"],)
if where == "WHERE ":
# no predicates, return an empty WHERE clause string
return ("", ())
else:
return (where, param_values)
return ("WHERE " + " AND ".join(wheres) if wheres else "", param_values)
langgraph/managed/few_shot.py
+4 -4
View File
@@ -66,15 +66,15 @@ class FewShotExamples(ManagedValue[Sequence[V]], Generic[V]):
return self.metadata_filter
def iter(self, score: int = 1) -> Iterator[V]:
for example in self.graph.checkpointer.search(
{"score": score, **self.metadata_filter_dict}, limit=self.k
for example in self.graph.checkpointer.list(
None, filter={"score": score, **self.metadata_filter_dict}, limit=self.k
):
with ChannelsManager(self.graph.channels, example.checkpoint) as channels:
yield read_channels(channels, self.graph.output_channels)
async def aiter(self, score: int = 1) -> AsyncIterator[V]:
async for example in self.graph.checkpointer.asearch(
{"score": score, **self.metadata_filter_dict}, limit=self.k
async for example in self.graph.checkpointer.alist(
None, filter={"score": score, **self.metadata_filter_dict}, limit=self.k
):
async with AsyncChannelsManager(
self.graph.channels, example.checkpoint
langgraph/pregel/__init__.py
+16 -4
View File
@@ -5,11 +5,13 @@ import concurrent.futures
import time
from collections import defaultdict, deque
from functools import partial
from inspect import signature
from typing import (
Any,
AsyncIterator,
Awaitable,
Callable,
Dict,
Iterator,
Literal,
Mapping,
@@ -410,15 +412,20 @@ class Pregel(
self,
config: RunnableConfig,
*,
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[StateSnapshot]:
"""Get the history of the state of the graph."""
if not self.checkpointer:
raise ValueError("No checkpointer set")
if (
filter is not None
and signature(self.checkpointer.list).parameters.get("filter") is None
):
raise ValueError("Checkpointer does not support filtering")
for config, checkpoint, metadata, parent_config in self.checkpointer.list(
config, before=before, limit=limit
config, before=before, limit=limit, filter=filter
):
with ChannelsManager(
self.channels, checkpoint
@@ -447,19 +454,24 @@ class Pregel(
self,
config: RunnableConfig,
*,
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[StateSnapshot]:
"""Get the history of the state of the graph."""
if not self.checkpointer:
raise ValueError("No checkpointer set")
if (
filter is not None
and signature(self.checkpointer.list).parameters.get("filter") is None
):
raise ValueError("Checkpointer does not support filtering")
async for (
config,
checkpoint,
metadata,
parent_config,
) in self.checkpointer.alist(config, before=before, limit=limit):
) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter):
async with AsyncChannelsManager(
self.channels, checkpoint
) as channels, AsyncManagedValuesManager(
tests/checkpoint/test_aiosqlite.py
+12 -4
View File
@@ -51,18 +51,26 @@ class TestAsyncSqliteSaver:
query_4: CheckpointMetadata = {"source": "update", "step": 1} # no match
async with self.sqlite_saver as sqlite_saver:
search_results_1 = [c async for c in sqlite_saver.asearch(query_1)]
search_results_1 = [
c async for c in sqlite_saver.alist(None, filter=query_1)
]
assert len(search_results_1) == 1
assert search_results_1[0].metadata == self.metadata_1
search_results_2 = [c async for c in sqlite_saver.asearch(query_2)]
search_results_2 = [
c async for c in sqlite_saver.alist(None, filter=query_2)
]
assert len(search_results_2) == 1
assert search_results_2[0].metadata == self.metadata_2
search_results_3 = [c async for c in sqlite_saver.asearch(query_3)]
search_results_3 = [
c async for c in sqlite_saver.alist(None, filter=query_3)
]
assert len(search_results_3) == 2
search_results_4 = [c async for c in sqlite_saver.asearch(query_4)]
search_results_4 = [
c async for c in sqlite_saver.alist(None, filter=query_4)
]
assert len(search_results_4) == 0
# TODO: test before and limit params
tests/checkpoint/test_memory.py
+16 -8
View File
@@ -50,18 +50,18 @@ class TestMemorySaver:
query_3: CheckpointMetadata = {} # search by no keys, return all checkpoints
query_4: CheckpointMetadata = {"source": "update", "step": 1} # no match
search_results_1 = list(self.memory_saver.search(query_1))
search_results_1 = list(self.memory_saver.list(None, filter=query_1))
assert len(search_results_1) == 1
assert search_results_1[0].metadata == self.metadata_1
search_results_2 = list(self.memory_saver.search(query_2))
search_results_2 = list(self.memory_saver.list(None, filter=query_2))
assert len(search_results_2) == 1
assert search_results_2[0].metadata == self.metadata_2
search_results_3 = list(self.memory_saver.search(query_3))
search_results_3 = list(self.memory_saver.list(None, filter=query_3))
assert len(search_results_3) == 2
search_results_4 = list(self.memory_saver.search(query_4))
search_results_4 = list(self.memory_saver.list(None, filter=query_4))
assert len(search_results_4) == 0
# TODO: test before and limit params
@@ -81,16 +81,24 @@ class TestMemorySaver:
query_3: CheckpointMetadata = {} # search by no keys, return all checkpoints
query_4: CheckpointMetadata = {"source": "update", "step": 1} # no match
search_results_1 = [c async for c in self.memory_saver.asearch(query_1)]
search_results_1 = [
c async for c in self.memory_saver.alist(None, filter=query_1)
]
assert len(search_results_1) == 1
assert search_results_1[0].metadata == self.metadata_1
search_results_2 = [c async for c in self.memory_saver.asearch(query_2)]
search_results_2 = [
c async for c in self.memory_saver.alist(None, filter=query_2)
]
assert len(search_results_2) == 1
assert search_results_2[0].metadata == self.metadata_2
search_results_3 = [c async for c in self.memory_saver.asearch(query_3)]
search_results_3 = [
c async for c in self.memory_saver.alist(None, filter=query_3)
]
assert len(search_results_3) == 2
search_results_4 = [c async for c in self.memory_saver.asearch(query_4)]
search_results_4 = [
c async for c in self.memory_saver.alist(None, filter=query_4)
]
assert len(search_results_4) == 0
tests/checkpoint/test_sqlite.py
+23 -13
View File
@@ -56,40 +56,50 @@ class TestSqliteSaver:
query_3: CheckpointMetadata = {} # search by no keys, return all checkpoints
query_4: CheckpointMetadata = {"source": "update", "step": 1} # no match
search_results_1 = list(self.sqlite_saver.search(query_1))
search_results_1 = list(self.sqlite_saver.list(None, filter=query_1))
assert len(search_results_1) == 1
assert search_results_1[0].metadata == self.metadata_1
search_results_2 = list(self.sqlite_saver.search(query_2))
search_results_2 = list(self.sqlite_saver.list(None, filter=query_2))
assert len(search_results_2) == 1
assert search_results_2[0].metadata == self.metadata_2
search_results_3 = list(self.sqlite_saver.search(query_3))
search_results_3 = list(self.sqlite_saver.list(None, filter=query_3))
assert len(search_results_3) == 2
search_results_4 = list(self.sqlite_saver.search(query_4))
search_results_4 = list(self.sqlite_saver.list(None, filter=query_4))
assert len(search_results_4) == 0
# TODO: test before and limit params
def test_search_where(self):
# call method / assertions
expected_predicate_1 = "WHERE json_extract(CAST(metadata AS TEXT), '$.source') = ? AND json_extract(CAST(metadata AS TEXT), '$.step') = ? AND json_extract(CAST(metadata AS TEXT), '$.writes') = ? AND json_extract(CAST(metadata AS TEXT), '$.score') = ? AND thread_ts < ? "
expected_param_values_1 = ("input", 2, "{}", 1, "1")
assert search_where(self.metadata_1, self.config_1) == (
expected_predicate_1 = "WHERE json_extract(CAST(metadata AS TEXT), '$.source') = ? AND json_extract(CAST(metadata AS TEXT), '$.step') = ? AND json_extract(CAST(metadata AS TEXT), '$.writes') = ? AND json_extract(CAST(metadata AS TEXT), '$.score') = ? AND thread_ts < ?"
expected_param_values_1 = ["input", 2, "{}", 1, "1"]
assert search_where(None, self.metadata_1, self.config_1) == (
expected_predicate_1,
expected_param_values_1,
)
def test_metadata_predicate(self):
# call method / assertions
expected_predicate_1 = "json_extract(CAST(metadata AS TEXT), '$.source') = ? AND json_extract(CAST(metadata AS TEXT), '$.step') = ? AND json_extract(CAST(metadata AS TEXT), '$.writes') = ? AND json_extract(CAST(metadata AS TEXT), '$.score') = ? "
expected_predicate_2 = "json_extract(CAST(metadata AS TEXT), '$.source') = ? AND json_extract(CAST(metadata AS TEXT), '$.step') = ? AND json_extract(CAST(metadata AS TEXT), '$.writes') = ? AND json_extract(CAST(metadata AS TEXT), '$.score') IS ? "
expected_predicate_3 = ""
expected_predicate_1 = [
"json_extract(CAST(metadata AS TEXT), '$.source') = ?",
"json_extract(CAST(metadata AS TEXT), '$.step') = ?",
"json_extract(CAST(metadata AS TEXT), '$.writes') = ?",
"json_extract(CAST(metadata AS TEXT), '$.score') = ?",
]
expected_predicate_2 = [
"json_extract(CAST(metadata AS TEXT), '$.source') = ?",
"json_extract(CAST(metadata AS TEXT), '$.step') = ?",
"json_extract(CAST(metadata AS TEXT), '$.writes') = ?",
"json_extract(CAST(metadata AS TEXT), '$.score') IS ?",
]
expected_predicate_3 = []
expected_param_values_1 = ("input", 2, "{}", 1)
expected_param_values_2 = ("loop", 1, '{"foo":"bar"}', None)
expected_param_values_3 = ()
expected_param_values_1 = ["input", 2, "{}", 1]
expected_param_values_2 = ["loop", 1, '{"foo":"bar"}', None]
expected_param_values_3 = []
assert _metadata_predicate(self.metadata_1) == (
expected_predicate_1,
tests/test_pregel.py
+4 -4
View File
@@ -2870,7 +2870,7 @@ Some examples of past conversations:
metadata = chkpnt_tuple_1.metadata
# not needed in application code, only for testing
hiscored = list(saver.search({"score": 1}))
hiscored = list(saver.list(None, filter={"score": 1}))
assert hiscored == []
# mark as "good"
@@ -2878,7 +2878,7 @@ Some examples of past conversations:
saver.put(config, checkpoint, metadata)
# not needed in application code, only for testing
hiscored = list(saver.search({"score": 1}))
hiscored = list(saver.list(None, filter={"score": 1}))
assert len(hiscored) == 1
assert hiscored[0].checkpoint["channel_values"]["messages"] == first_messages
@@ -2921,14 +2921,14 @@ Some examples of past conversations:
metadata = chkpnt_tuple_2.metadata
# not needed in application code, only for testing
hiscored = list(saver.search({"score": 1}))
hiscored = list(saver.list(None, filter={"score": 1}))
assert len(hiscored) == 1
# mark as "good"
metadata["score"] = 1
saver.put(config, checkpoint, metadata)
hiscored = list(saver.search({"score": 1}))
hiscored = list(saver.list(None, filter={"score": 1}))
assert len(hiscored) == 2
assert app.invoke(
tests/test_pregel_async.py
+2 -2
View File
@@ -2643,14 +2643,14 @@ Some examples of past conversations:
metadata = chkpnt_tuple_1.metadata
# not needed in application code, only for testing
assert [c async for c in saver.asearch({"score": 1})] == []
assert [c async for c in saver.alist(None, filter={"score": 1})] == []
# mark as "good"
metadata["score"] = 1
await saver.aput(config, checkpoint, metadata)
# not needed in application code, only for testing
hiscored = [c async for c in saver.asearch({"score": 1})]
hiscored = [c async for c in saver.alist(None, filter={"score": 1})]
assert len(hiscored) == 1
assert hiscored[0].checkpoint["channel_values"]["messages"] == first_messages