fix msg routing bug, use message class

metagpt/roles/di/role_zero.py

@@ -13,7 +13,7 @@ from metagpt.environment.mgx.mgx_env import MGXEnv
 from metagpt.logs import logger
 from metagpt.prompts.di.role_zero import CMD_PROMPT, ROLE_INSTRUCTION
 from metagpt.roles import Role
-from metagpt.schema import Message
+from metagpt.schema import AIMessage, Message, UserMessage
 from metagpt.strategy.experience_retriever import DummyExpRetriever, ExpRetriever
 from metagpt.strategy.planner import Planner
 from metagpt.tools.tool_recommend import BM25ToolRecommender, ToolRecommender
@@ -117,11 +117,11 @@ class RoleZero(Role):
             available_commands=tool_info,
             instruction=self.instruction.strip(),
         )
-        context = self.llm.format_msg(self.rc.memory.get(self.memory_k) + [Message(content=prompt, role="user")])
+        context = self.llm.format_msg(self.rc.memory.get(self.memory_k) + [UserMessage(content=prompt)])
         print(*context, sep="\n" + "*" * 5 + "\n")
         async with ThoughtReporter():
             self.command_rsp = await self.llm.aask(context, system_msgs=self.system_msg)
-        self.rc.memory.add(Message(content=self.command_rsp, role="assistant"))
+        self.rc.memory.add(AIMessage(content=self.command_rsp))
 
         return True
 
@@ -134,21 +134,20 @@ class RoleZero(Role):
         except Exception as e:
             tb = traceback.format_exc()
             print(tb)
-            error_msg = Message(content=str(e), role="user")
+            error_msg = UserMessage(content=str(e))
             self.rc.memory.add(error_msg)
             return error_msg
         outputs = await self._run_commands(commands)
-        self.rc.memory.add(Message(content=outputs, role="user"))
-        return Message(
+        self.rc.memory.add(UserMessage(content=outputs))
+        return AIMessage(
             content=f"Complete run with outputs: {outputs}",
-            role="assistant",
-            sent_from=self._setting,
+            sent_from=self.name,
             cause_by=RunCommand,
         )
 
     async def _react(self) -> Message:
         actions_taken = 0
-        rsp = Message(content="No actions taken yet", cause_by=Action)  # will be overwritten after Role _act
+        rsp = AIMessage(content="No actions taken yet", cause_by=Action)  # will be overwritten after Role _act
         while actions_taken < self.rc.max_react_loop:
             # NOTE: difference here, keep observing within react
             await self._observe()

metagpt/roles/di/team_leader.py

@@ -2,13 +2,14 @@ from __future__ import annotations
 
 from pydantic import model_validator
 
+from metagpt.actions.di.run_command import RunCommand
 from metagpt.prompts.di.team_leader import (
     FINISH_CURRENT_TASK_CMD,
     SYSTEM_PROMPT,
     TL_INSTRUCTION,
 )
 from metagpt.roles.di.role_zero import RoleZero
-from metagpt.schema import Message
+from metagpt.schema import AIMessage, Message, UserMessage
 from metagpt.strategy.experience_retriever import ExpRetriever, SimpleExpRetriever
 from metagpt.tools.tool_registry import register_tool
 
@@ -49,7 +50,7 @@ class TeamLeader(RoleZero):
         self.set_instruction()
         return await super()._think()
 
-    def publish_message(self, msg, send_to="no one"):
+    def publish_message(self, msg: Message, send_to="no one"):
         """Overwrite Role.publish_message, send to no one if called within Role.run, send to the specified role if called dynamically."""
         if not msg:
             return
@@ -64,8 +65,11 @@ class TeamLeader(RoleZero):
         Publish a message to a team member, use member name to fill send_to args. You may copy the full original content or add additional information from upstream. This will make team members start their work.
         DONT omit any necessary info such as path, link, environment, programming language, framework, requirement, constraint from original content to team members because you are their sole info source.
         """
-        self.publish_message(Message(content=content), send_to=send_to)
+        # Specify the outer send_to to overwrite the default "no one" value. Use UserMessage because message from self is like a user request for others.
+        self.publish_message(
+            UserMessage(content=content, sent_from=self.name, send_to=send_to, cause_by=RunCommand), send_to=send_to
+        )
 
     def finish_current_task(self):
         self.planner.plan.finish_current_task()
-        self.rc.memory.add(Message(content=FINISH_CURRENT_TASK_CMD, role="assistant"))
+        self.rc.memory.add(AIMessage(content=FINISH_CURRENT_TASK_CMD))