fix tests

shihanwan · Oct 16, 2024 · 38af4c1 · 38af4c1
1 parent e635732
commit 38af4c1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/memonto/core/retain.py b/memonto/core/retain.py
@@ -158,10 +158,11 @@ def find_relevant_memories(
         relevant_memory = str(data.serialize(format="turtle"))
     else:
         relevant_memory = str(vector_store.search(message=message, id=id, k=3))
-    
+
     logger.debug(f"relevant_memory\n{relevant_memory}\n")
     return relevant_memory
 
+
 def save_memory(
     ontology: Graph,
     namespaces: dict[str, Namespace],
@@ -182,7 +183,6 @@ def save_memory(
         id=id,
         ephemeral=ephemeral,
     )
-    print(relevant_memory)
 
     script = llm.prompt(
         prompt_name="commit_to_memory",

diff --git a/tests/core/test_retain.py b/tests/core/test_retain.py
@@ -66,6 +66,7 @@ def test_commit_memory(
         ontology=ANY,
         user_message=user_query,
         updated_memory="",
+        relevant_memory=ANY,
     )
 
     assert mock_llm.prompt.call_count == 1
@@ -103,6 +104,7 @@ def test_commit_memory_with_exception(
         ontology=ANY,
         user_message=user_query,
         updated_memory="",
+        relevant_memory=ANY,
     )
 
     ctmeh_prompt = call(
@@ -146,6 +148,7 @@ def test_commit_memory_auto_expand(
         ontology=ANY,
         user_message=user_query,
         updated_memory="",
+        relevant_memory=ANY,
     )
 
     eo_prompt = call(