diff --git a/src/agentscope/plan/_plan_notebook.py b/src/agentscope/plan/_plan_notebook.py index e4c1afc18..efbef1832 100644 --- a/src/agentscope/plan/_plan_notebook.py +++ b/src/agentscope/plan/_plan_notebook.py @@ -221,7 +221,7 @@ def __init__( # Register the current_plan state for state management self.register_state( "current_plan", - custom_to_json=lambda _: _.model_dump_json() if _ else None, + custom_to_json=lambda _: _.model_dump() if _ else None, custom_from_json=lambda _: Plan.model_validate(_) if _ else None, ) diff --git a/tests/plan_test.py b/tests/plan_test.py index 8994d8ca6..98facdd39 100644 --- a/tests/plan_test.py +++ b/tests/plan_test.py @@ -1,7 +1,11 @@ # -*- coding: utf-8 -*- """The plan module related tests.""" +import os from unittest import IsolatedAsyncioTestCase +from agentscope.agent import ReActAgent +from agentscope.formatter import DashScopeChatFormatter +from agentscope.model import DashScopeChatModel from agentscope.plan import SubTask, Plan, PlanNotebook @@ -269,3 +273,80 @@ async def test_plan_subtasks(self) -> None: plan_notebook.current_plan.subtasks[1].state, "in_progress", ) + + async def test_serialization(self) -> None: + """Test the serialization and deserialization of plan and subtask.""" + plan_notebook = PlanNotebook() + agent = ReActAgent( + name="Friday", + sys_prompt="You are a helpful assistant named Friday. ", + model=DashScopeChatModel( + model_name="qwen-max", + api_key=os.environ.get("DASH_API_KEY"), + ), + formatter=DashScopeChatFormatter(), + plan_notebook=plan_notebook, + ) + + await plan_notebook.create_plan( + name="text", + description="abc", + expected_outcome="edf", + subtasks=[ + SubTask( + name="1", + description="1", + expected_outcome="1", + ), + SubTask( + name="2", + description="2", + expected_outcome="2", + ), + ], + ) + + self.assertIsNotNone(plan_notebook.current_plan) + + state = agent.state_dict() + subtasks = plan_notebook.current_plan.subtasks + self.assertDictEqual( + state, + { + "memory": {"content": []}, + "toolkit": {"active_groups": []}, + "plan_notebook": { + "storage": {}, + "current_plan": { + "name": "text", + "description": "abc", + "expected_outcome": "edf", + "subtasks": [ + { + "name": "1", + "description": "1", + "expected_outcome": "1", + "created_at": subtasks[0].created_at, + }, + { + "name": "2", + "description": "2", + "expected_outcome": "2", + "created_at": subtasks[1].created_at, + }, + ], + }, + }, + "_reasoning_hint_msgs": {"content": []}, + "name": "Friday", + "_sys_prompt": "You are a helpful assistant named Friday. ", + }, + ) + plan_notebook.current_plan = None + self.assertIsNone( + agent.plan_notebook.current_plan, + ) + agent.load_state_dict(state) + self.assertIsNotNone( + agent.plan_notebook.current_plan, + )