From d366bacf033d96de3f8e7782d74404de70e1f4c3 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 5 Jun 2024 17:14:35 -0700 Subject: [PATCH] Expose filter arg in get_state_history - Combine list and search methods in Checkpointer --- langgraph/checkpoint/aiosqlite.py | 96 +++--------------- langgraph/checkpoint/base.py | 25 +---- langgraph/checkpoint/memory.py | 111 +++++---------------- langgraph/checkpoint/sqlite.py | 154 +++++++---------------------- langgraph/managed/few_shot.py | 8 +- langgraph/pregel/__init__.py | 20 +++- tests/checkpoint/test_aiosqlite.py | 16 ++- tests/checkpoint/test_memory.py | 24 +++-- tests/checkpoint/test_sqlite.py | 36 ++++--- tests/test_pregel.py | 8 +- tests/test_pregel_async.py | 4 +- 11 files changed, 154 insertions(+), 348 deletions(-) diff --git a/langgraph/checkpoint/aiosqlite.py b/langgraph/checkpoint/aiosqlite.py index 1e9734e7a..73341a560 100644 --- a/langgraph/checkpoint/aiosqlite.py +++ b/langgraph/checkpoint/aiosqlite.py @@ -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}}, diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 7fb88eedc..a639bc336 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -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, diff --git a/langgraph/checkpoint/memory.py b/langgraph/checkpoint/memory.py index 2f592cec5..d94eb3d64 100644 --- a/langgraph/checkpoint/memory.py +++ b/langgraph/checkpoint/memory.py @@ -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, diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index 23c054912..3836a041f 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/langgraph/checkpoint/sqlite.py @@ -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) diff --git a/langgraph/managed/few_shot.py b/langgraph/managed/few_shot.py index a96710b55..9d46739eb 100644 --- a/langgraph/managed/few_shot.py +++ b/langgraph/managed/few_shot.py @@ -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 diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index d42839db7..74b404f7b 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -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( diff --git a/tests/checkpoint/test_aiosqlite.py b/tests/checkpoint/test_aiosqlite.py index fe05296ea..99e24b722 100644 --- a/tests/checkpoint/test_aiosqlite.py +++ b/tests/checkpoint/test_aiosqlite.py @@ -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 diff --git a/tests/checkpoint/test_memory.py b/tests/checkpoint/test_memory.py index 32de877a4..e41e53528 100644 --- a/tests/checkpoint/test_memory.py +++ b/tests/checkpoint/test_memory.py @@ -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 diff --git a/tests/checkpoint/test_sqlite.py b/tests/checkpoint/test_sqlite.py index bc678d7b9..3753e5fc6 100644 --- a/tests/checkpoint/test_sqlite.py +++ b/tests/checkpoint/test_sqlite.py @@ -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, diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 18dc2759c..24566598a 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -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( diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index dabdf964f..51a024fe3 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -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