Expose filter arg in get_state_history

- 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
+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