From 831ddb1736ecee40106b1503ae4cd2e8da2349ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=8E=98=E6=9D=83=20=E9=A9=AC?= <mashenquan@fuzhi.ai>
Date: Fri, 19 Jan 2024 17:24:05 +0800
Subject: [PATCH] feat: + composition, aggregation relationship

---
 metagpt/actions/rebuild_sequence_view.py | 45 +++++++++++++++++++-----
 1 file changed, 36 insertions(+), 9 deletions(-)

diff --git a/metagpt/actions/rebuild_sequence_view.py b/metagpt/actions/rebuild_sequence_view.py
index 8cb656c09..d4cba63ca 100644
--- a/metagpt/actions/rebuild_sequence_view.py
+++ b/metagpt/actions/rebuild_sequence_view.py
@@ -18,7 +18,7 @@ from tenacity import retry, stop_after_attempt, wait_random_exponential
 
 from metagpt.actions import Action
 from metagpt.config2 import config
-from metagpt.const import GENERALIZATION, GRAPH_REPO_FILE_REPO
+from metagpt.const import AGGREGATION, COMPOSITION, GENERALIZATION, GRAPH_REPO_FILE_REPO
 from metagpt.logs import logger
 from metagpt.schema import ClassView
 from metagpt.utils.common import aread, general_after_log, list_files, split_namespace
@@ -164,14 +164,28 @@ class RebuildSequenceView(Action):
     )
     async def _merge_sequence_view(self, sequence_view, ns_class_name) -> str:
         class_view_part = "```mermaid\n"
-        class_view_rows = await self.graph_db.select(subject=ns_class_name, predicate=GraphKeyword.HAS_CLASS_VIEW)
-        class_view_part += (
-            ClassView.model_validate_json(class_view_rows[0].object_).get_mermaid() if class_view_rows else ""
+        # add class
+        class_view_part += await self._class_view_to_mermaid(ns_class_name)
+        # add aggregation relationship
+        _, me = split_namespace(ns_class_name)
+        rows = await self.graph_db.select(
+            subject=ns_class_name, predicate=GraphKeyword.IS + AGGREGATION + GraphKeyword.OF
         )
-        _, name = split_namespace(class_view_rows[0].subject)
-        class_desc_rows = await self.graph_db.select(subject=ns_class_name, predicate=GraphKeyword.HAS_CLASS_DESC)
-        for r in self._desc_to_note(class_desc_rows[0].object_):
-            class_view_part += f'\n\tnote for {name} "{r}"'
+        aggregation = [r.object_ for r in rows]
+        for ns_aggr_name in aggregation:
+            _, name = split_namespace(ns_aggr_name)
+            class_view_part += f"\n\t{me} *-- {name}"
+            class_view_part += await self._class_view_to_mermaid(ns_aggr_name)
+        # add composition relationship
+        rows = await self.graph_db.select(
+            subject=ns_class_name, predicate=GraphKeyword.IS + COMPOSITION + GraphKeyword.OF
+        )
+        compositions = [r.object_ for r in rows]
+        for ns_comp_name in compositions:
+            _, name = split_namespace(ns_comp_name)
+            class_view_part += f"\n\t{me} *-- {name}"
+            class_view_part += await self._class_view_to_mermaid(ns_comp_name)
+
         class_view_part += "\n```"
 
         contents = [
@@ -183,7 +197,7 @@ class RebuildSequenceView(Action):
             msg=prompt,
             system_msgs=[
                 "You are a tool to merge Mermaid class view information into the Mermaid sequence view.",
-                'Append as much information as possible from the "Mermaid Class View" and "Class Description" to the sequence diagram.',
+                'Append as much information as possible from the "Mermaid Class View" to the sequence diagram.',
                 'Return a markdown JSON format with a "sequence_diagram" key containing the merged Mermaid sequence view, a "reason" key explaining what information have been merged and why.',
             ],
         )
@@ -234,3 +248,16 @@ class RebuildSequenceView(Action):
             return []
         m = json.loads(json_str)
         return [s.replace('"', '\\"').replace("\n", "\\n") for s in m.values()]
+
+    async def _class_view_to_mermaid(self, ns_class_name):
+        class_view_rows = await self.graph_db.select(subject=ns_class_name, predicate=GraphKeyword.HAS_CLASS_VIEW)
+        result = ClassView.model_validate_json(class_view_rows[0].object_).get_mermaid() if class_view_rows else ""
+        _, name = split_namespace(ns_class_name)
+        class_desc_rows = await self.graph_db.select(subject=ns_class_name, predicate=GraphKeyword.HAS_CLASS_DESC)
+        if not class_desc_rows:
+            # Haven't been parsed before.
+            await self._parse_class_description(ns_class_name=ns_class_name)
+            class_desc_rows = await self.graph_db.select(subject=ns_class_name, predicate=GraphKeyword.HAS_CLASS_DESC)
+        for r in self._desc_to_note(class_desc_rows[0].object_):
+            result += f'\n\tnote for {name} "{r}"'
+        return result