From 1f6af59e9dcb099bba63ff8f4fcdc25d268ca6de Mon Sep 17 00:00:00 2001 From: cpacker Date: Tue, 2 Jan 2024 18:20:54 -0800 Subject: [PATCH] added a new test that makes sure the base functions are working --- tests/test_base_functions.py | 50 ++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) create mode 100644 tests/test_base_functions.py diff --git a/tests/test_base_functions.py b/tests/test_base_functions.py new file mode 100644 index 00000000..d2a3f695 --- /dev/null +++ b/tests/test_base_functions.py @@ -0,0 +1,50 @@ +from memgpt import MemGPT +from memgpt import constants +import memgpt.functions.function_sets.base as base_functions + +from .utils import wipe_config + + +# test_agent_id = "test_agent" +client = None +agent_obj = None + + +def create_test_agent(): + """Create a test agent that we can call functions on""" + wipe_config() + global client + client = MemGPT(quickstart="memgpt_hosted") + + agent_id = client.create_agent( + agent_config={ + # "name": test_agent_id, + "persona": constants.DEFAULT_PERSONA, + "human": constants.DEFAULT_HUMAN, + } + ) + + global agent_obj + agent_obj = client.server._get_or_load_agent(user_id="NULL", agent_id=agent_id) + + +def test_archival(): + global agent_obj + if agent_obj is None: + create_test_agent() + + base_functions.archival_memory_insert(agent_obj, "banana") + + base_functions.archival_memory_search(agent_obj, "banana") + base_functions.archival_memory_search(agent_obj, "banana", page=0) + + +def test_recall(): + global agent_obj + if agent_obj is None: + create_test_agent() + + base_functions.conversation_search(agent_obj, "banana") + base_functions.conversation_search(agent_obj, "banana", page=0) + + base_functions.conversation_search_date(agent_obj, start_date="2022-01-01", end_date="2022-01-02")