From edce4ac47afed589b81bd856b5cf3752ccd41329 Mon Sep 17 00:00:00 2001 From: better629 Date: Fri, 29 Dec 2023 03:10:23 +0800 Subject: [PATCH] fix memory unittest --- metagpt/memory/longterm_memory.py | 5 +++-- tests/metagpt/memory/test_longterm_memory.py | 4 ++++ tests/metagpt/memory/test_memory_storage.py | 4 ++-- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/metagpt/memory/longterm_memory.py b/metagpt/memory/longterm_memory.py index 8da6ed84a..b54653970 100644 --- a/metagpt/memory/longterm_memory.py +++ b/metagpt/memory/longterm_memory.py @@ -12,6 +12,7 @@ from pydantic import ConfigDict, Field from metagpt.logs import logger from metagpt.memory import Memory from metagpt.memory.memory_storage import MemoryStorage +from metagpt.roles.role import RoleContext from metagpt.schema import Message @@ -25,10 +26,10 @@ class LongTermMemory(Memory): model_config = ConfigDict(arbitrary_types_allowed=True) memory_storage: MemoryStorage = Field(default_factory=MemoryStorage) - rc: Optional["RoleContext"] = None + rc: Optional[RoleContext] = None msg_from_recover: bool = False - def recover_memory(self, role_id: str, rc: "RoleContext"): + def recover_memory(self, role_id: str, rc: RoleContext): messages = self.memory_storage.recover_memory(role_id) self.rc = rc if not self.memory_storage.is_initialized: diff --git a/tests/metagpt/memory/test_longterm_memory.py b/tests/metagpt/memory/test_longterm_memory.py index ac33552b3..c915a6610 100644 --- a/tests/metagpt/memory/test_longterm_memory.py +++ b/tests/metagpt/memory/test_longterm_memory.py @@ -20,6 +20,10 @@ def test_ltm_search(): assert len(CONFIG.openai_api_key) > 20 role_id = "UTUserLtm(Product Manager)" + from metagpt.environment import Environment + + Environment + RoleContext.model_rebuild() rc = RoleContext(watch={"metagpt.actions.add_requirement.UserRequirement"}) ltm = LongTermMemory() ltm.recover_memory(role_id, rc) diff --git a/tests/metagpt/memory/test_memory_storage.py b/tests/metagpt/memory/test_memory_storage.py index f1cc12aac..0eb1069d5 100644 --- a/tests/metagpt/memory/test_memory_storage.py +++ b/tests/metagpt/memory/test_memory_storage.py @@ -24,7 +24,7 @@ def test_idea_message(): role_id = "UTUser1(Product Manager)" message = Message(role="User", content=idea, cause_by=UserRequirement) - shutil.rmtree(Path(DATA_PATH / f"role_mem/{role_id}/")) + shutil.rmtree(Path(DATA_PATH / f"role_mem/{role_id}/"), ignore_errors=True) memory_storage: MemoryStorage = MemoryStorage() messages = memory_storage.recover_memory(role_id) @@ -58,7 +58,7 @@ def test_actionout_message(): content=content, instruct_content=ic_obj(**out_data), role="user", cause_by=WritePRD ) # WritePRD as test action - shutil.rmtree(Path(DATA_PATH / f"role_mem/{role_id}/")) + shutil.rmtree(Path(DATA_PATH / f"role_mem/{role_id}/"), ignore_errors=True) memory_storage: MemoryStorage = MemoryStorage() messages = memory_storage.recover_memory(role_id)