From e2e00beb755bb10c73460cad2f19944567cbd4ea Mon Sep 17 00:00:00 2001
From: better629 <webdesign_xmu2012@163.com>
Date: Tue, 9 Jan 2024 15:40:42 +0800
Subject: [PATCH] make instruct_content support any inherited basemodel
 ser&deser

---
 metagpt/schema.py                             | 25 ++++---
 .../serialize_deserialize/test_schema.py      | 68 +++++++++++++++----
 .../test_serdeser_base.py                     | 10 +--
 3 files changed, 77 insertions(+), 26 deletions(-)

diff --git a/metagpt/schema.py b/metagpt/schema.py
index a557951c7..7d1c2b539 100644
--- a/metagpt/schema.py
+++ b/metagpt/schema.py
@@ -182,12 +182,16 @@ class Message(BaseModel):
     @field_validator("instruct_content", mode="before")
     @classmethod
     def check_instruct_content(cls, ic: Any) -> BaseModel:
-        if ic and not isinstance(ic, BaseModel) and "class" in ic:
-            # compatible with custom-defined ActionOutput
-            mapping = actionoutput_str_to_mapping(ic["mapping"])
-
-            actionnode_class = import_class("ActionNode", "metagpt.actions.action_node")  # avoid circular import
-            ic_obj = actionnode_class.create_model_class(class_name=ic["class"], mapping=mapping)
+        if ic and isinstance(ic, dict) and "class" in ic:
+            if "mapping" in ic:
+                # compatible with custom-defined ActionOutput
+                mapping = actionoutput_str_to_mapping(ic["mapping"])
+                actionnode_class = import_class("ActionNode", "metagpt.actions.action_node")  # avoid circular import
+                ic_obj = actionnode_class.create_model_class(class_name=ic["class"], mapping=mapping)
+            elif "module" in ic:
+                ic_obj = import_class(ic["class"], ic["module"])
+            else:
+                raise KeyError("missing required key to init Message.instruct_content from dict")
             ic = ic_obj(**ic["value"])
         return ic
 
@@ -212,13 +216,16 @@ class Message(BaseModel):
         if ic:
             # compatible with custom-defined ActionOutput
             schema = ic.model_json_schema()
-            # `Documents` contain definitions
-            if "definitions" not in schema:
-                # TODO refine with nested BaseModel
+            ic_type = str(type(ic))
+            if "<class 'metagpt.actions.action_node" in ic_type:
+                # instruct_content from AutoNode.create_model_class, for now, it's single level structure.
                 mapping = actionoutout_schema_to_mapping(schema)
                 mapping = actionoutput_mapping_to_str(mapping)
 
                 ic_dict = {"class": schema["title"], "mapping": mapping, "value": ic.model_dump()}
+            else:
+                # due to instruct_content can be assigned by subclasses of BaseModel
+                ic_dict = {"class": schema["title"], "module": ic.__module__, "value": ic.model_dump()}
         return ic_dict
 
     def __init__(self, content: str = "", **data: Any):
diff --git a/tests/metagpt/serialize_deserialize/test_schema.py b/tests/metagpt/serialize_deserialize/test_schema.py
index e793079f0..b55b82088 100644
--- a/tests/metagpt/serialize_deserialize/test_schema.py
+++ b/tests/metagpt/serialize_deserialize/test_schema.py
@@ -1,10 +1,11 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 # @Desc   : unittest of schema ser&deser
+import pytest
 
 from metagpt.actions.action_node import ActionNode
 from metagpt.actions.write_code import WriteCode
-from metagpt.schema import Document, Documents, Message
+from metagpt.schema import CodingContext, Document, Documents, Message, TestingContext
 from metagpt.utils.common import any_to_str
 from tests.metagpt.serialize_deserialize.test_serdeser_base import (
     MockICMessage,
@@ -12,12 +13,16 @@ from tests.metagpt.serialize_deserialize.test_serdeser_base import (
 )
 
 
-def test_message_serdeser():
+def test_message_serdeser_from_create_model():
+    with pytest.raises(KeyError):
+        _ = Message(content="code", instruct_content={"class": "test", "key": "value"})
+
     out_mapping = {"field3": (str, ...), "field4": (list[str], ...)}
     out_data = {"field3": "field3 value3", "field4": ["field4 value1", "field4 value2"]}
     ic_obj = ActionNode.create_model_class("code", out_mapping)
+    ic_inst = ic_obj(**out_data)
 
-    message = Message(content="code", instruct_content=ic_obj(**out_data), role="engineer", cause_by=WriteCode)
+    message = Message(content="code", instruct_content=ic_inst, role="engineer", cause_by=WriteCode)
     ser_data = message.model_dump()
     assert ser_data["cause_by"] == "metagpt.actions.write_code.WriteCode"
     assert ser_data["instruct_content"]["class"] == "code"
@@ -25,28 +30,67 @@ def test_message_serdeser():
     new_message = Message(**ser_data)
     assert new_message.cause_by == any_to_str(WriteCode)
     assert new_message.cause_by in [any_to_str(WriteCode)]
+
     assert new_message.instruct_content != ic_obj(**out_data)  # TODO find why `!=`
+    assert new_message.instruct_content != ic_inst
     assert new_message.instruct_content.model_dump() == ic_obj(**out_data).model_dump()
 
-    message = Message(content="test_ic", instruct_content=MockICMessage())
+    mock_msg = MockMessage()
+    message = Message(content="test_ic", instruct_content=mock_msg)
     ser_data = message.model_dump()
     new_message = Message(**ser_data)
-    assert new_message.instruct_content != MockICMessage()  # TODO
-
-    message = Message(content="test_documents", instruct_content=Documents(docs={"doc1": Document(content="test doc")}))
-    ser_data = message.model_dump()
-    assert "class" in ser_data["instruct_content"]
+    assert new_message.instruct_content == mock_msg
 
 
 def test_message_without_postprocess():
-    """to explain `instruct_content` should be postprocessed"""
+    """to explain `instruct_content` from `create_model_class` should be postprocessed"""
     out_mapping = {"field1": (list[str], ...)}
     out_data = {"field1": ["field1 value1", "field1 value2"]}
     ic_obj = ActionNode.create_model_class("code", out_mapping)
-    message = MockMessage(content="code", instruct_content=ic_obj(**out_data))
+    message = MockICMessage(content="code", instruct_content=ic_obj(**out_data))
     ser_data = message.model_dump()
     assert ser_data["instruct_content"] == {}
 
     ser_data["instruct_content"] = None
-    new_message = MockMessage(**ser_data)
+    new_message = MockICMessage(**ser_data)
     assert new_message.instruct_content != ic_obj(**out_data)
+
+
+def test_message_serdeser_from_basecontext():
+    doc_msg = Message(content="test_document", instruct_content=Document(content="test doc"))
+    ser_data = doc_msg.model_dump()
+    assert ser_data["instruct_content"]["value"]["content"] == "test doc"
+    assert ser_data["instruct_content"]["value"]["filename"] == ""
+
+    docs_msg = Message(
+        content="test_documents", instruct_content=Documents(docs={"doc1": Document(content="test doc")})
+    )
+    ser_data = docs_msg.model_dump()
+    assert ser_data["instruct_content"]["class"] == "Documents"
+    assert ser_data["instruct_content"]["value"]["docs"]["doc1"]["content"] == "test doc"
+    assert ser_data["instruct_content"]["value"]["docs"]["doc1"]["filename"] == ""
+
+    code_ctxt = CodingContext(
+        filename="game.py",
+        design_doc=Document(root_path="docs/system_design", filename="xx.json", content="xxx"),
+        task_doc=Document(root_path="docs/tasks", filename="xx.json", content="xxx"),
+        code_doc=Document(root_path="xxx", filename="game.py", content="xxx"),
+    )
+    code_ctxt_msg = Message(content="coding_context", instruct_content=code_ctxt)
+    ser_data = code_ctxt_msg.model_dump()
+    assert ser_data["instruct_content"]["class"] == "CodingContext"
+
+    new_code_ctxt_msg = Message(**ser_data)
+    assert new_code_ctxt_msg.instruct_content == code_ctxt
+    assert new_code_ctxt_msg.instruct_content.code_doc.filename == "game.py"
+
+    testing_ctxt = TestingContext(
+        filename="test.py",
+        code_doc=Document(root_path="xxx", filename="game.py", content="xxx"),
+        test_doc=Document(root_path="docs/tests", filename="test.py", content="xxx"),
+    )
+    testing_ctxt_msg = Message(content="testing_context", instruct_content=testing_ctxt)
+    ser_data = testing_ctxt_msg.model_dump()
+    new_testing_ctxt_msg = Message(**ser_data)
+    assert new_testing_ctxt_msg.instruct_content == testing_ctxt
+    assert new_testing_ctxt_msg.instruct_content.test_doc.filename == "test.py"
diff --git a/tests/metagpt/serialize_deserialize/test_serdeser_base.py b/tests/metagpt/serialize_deserialize/test_serdeser_base.py
index daa46c99c..ddb47a3e2 100644
--- a/tests/metagpt/serialize_deserialize/test_serdeser_base.py
+++ b/tests/metagpt/serialize_deserialize/test_serdeser_base.py
@@ -16,14 +16,14 @@ from metagpt.roles.role import Role, RoleReactMode
 serdeser_path = Path(__file__).absolute().parent.joinpath("..", "..", "data", "serdeser_storage")
 
 
-class MockICMessage(BaseModel):
-    content: str = "test_ic"
-
-
 class MockMessage(BaseModel):
+    content: str = "test_msg"
+
+
+class MockICMessage(BaseModel):
     """to test normal dict without postprocess"""
 
-    content: str = ""
+    content: str = "test_ic_msg"
     instruct_content: Optional[BaseModel] = Field(default=None)