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))

tests/metagpt/roles/di/run_engineer2.py

@@ -3,9 +3,7 @@ import asyncio
 from metagpt.roles.di.engineer2 import Engineer2
 
 DESIGN_DOC_2048 = '{"Implementation approach":"We will use the Pygame library to implement the 2048 game logic and user interface. Pygame is a set of Python modules designed for writing video games, which will help us create a responsive and visually appealing UI. For the mobile responsiveness, we will ensure that the game scales appropriately on different screen sizes. We will also use the Pygame GUI library to create buttons for restarting the game and choosing difficulty levels.","File list":["main.py","game.py","ui.py"],"Data structures and interfaces":"\\nclassDiagram\\n    class Game {\\n        -grid: list[list[int]]\\n        -score: int\\n        +__init__()\\n        +move(direction: str) bool\\n        +merge() bool\\n        +spawn_tile() None\\n        +is_game_over() bool\\n        +reset() None\\n    }\\n    class UI {\\n        -game: Game\\n        +__init__(game: Game)\\n        +draw_grid() None\\n        +draw_score() None\\n        +draw_buttons() None\\n        +handle_input() None\\n    }\\n    class Main {\\n        -ui: UI\\n        +main() None\\n    }\\n    Main --> UI\\n    UI --> Game\\n","Program call flow":"\\nsequenceDiagram\\n    participant M as Main\\n    participant U as UI\\n    participant G as Game\\n    M->>U: __init__(game)\\n    U->>G: __init__()\\n    M->>U: draw_grid()\\n    U->>G: move(direction)\\n    G-->>U: return bool\\n    U->>G: merge()\\n    G-->>U: return bool\\n    U->>G: spawn_tile()\\n    G-->>U: return None\\n    U->>G: is_game_over()\\n    G-->>U: return bool\\n    U->>G: reset()\\n    G-->>U: return None\\n    M->>U: draw_score()\\n    M->>U: draw_buttons()\\n    M->>U: handle_input()\\n","Anything UNCLEAR":"Clarification needed on the specific design elements for the UI to ensure it meets the \'beautiful\' requirement. Additionally, we need to confirm the exact difficulty levels and how they should affect the game mechanics."}'
-
 TASK_DOC_2048 = '{"Required Python packages":["pygame==2.0.1","pygame_gui==0.5.7"],"Required Other language third-party packages":["No third-party dependencies required"],"Logic Analysis":[["game.py","Contains Game class with methods: __init__, move, merge, spawn_tile, is_game_over, reset"],["ui.py","Contains UI class with methods: __init__, draw_grid, draw_score, draw_buttons, handle_input"],["main.py","Contains Main class with method: main, initializes UI and Game"]],"Task list":["game.py","ui.py","main.py"],"Full API spec":"","Shared Knowledge":"`game.py` contains core game logic and state management. `ui.py` handles all user interface elements and interactions. `main.py` serves as the entry point to initialize and run the game.","Anything UNCLEAR":"Clarification needed on the specific design elements for the UI to ensure it meets the \'beautiful\' requirement. Additionally, we need to confirm the exact difficulty levels and how they should affect the game mechanics."}'
-
 DESIGN_DOC_SNAKE = """
 {
     "Implementation approach": "We will use the Pygame library to create the CLI-based snake game. Pygame is a set of Python modules designed for writing video games, which will help us handle graphics, sound, and input. The game will be structured into different modules to handle the main game loop, snake movement, food generation, collision detection, and user interface. We will ensure the game is engaging and responsive by optimizing the game loop and input handling. The score display and different speed levels will be implemented to enhance the user experience.",