From 401fe3aabe00f3d8d96bf509b656fd848bac4e44 Mon Sep 17 00:00:00 2001
From: DavdGao <gaodawei.gdw@alibaba-inc.com>
Date: Tue, 2 Jul 2024 10:29:59 +0800
Subject: [PATCH] Unify the typing of the agents' reply function and update
 their docstrings. (#314)

---
 .../en/source/tutorial/104-usecase.md         |  6 +++++-
 .../en/source/tutorial/201-agent.md           | 10 +++++-----
 .../en/source/tutorial/203-parser.md          |  2 +-
 .../en/source/tutorial/204-service.md         |  5 ++++-
 .../en/source/tutorial/209-prompt_opt.md      |  4 +++-
 .../zh_CN/source/tutorial/104-usecase.md      |  7 ++++++-
 .../zh_CN/source/tutorial/201-agent.md        |  8 ++++----
 .../zh_CN/source/tutorial/203-parser.md       |  2 +-
 .../zh_CN/source/tutorial/204-service.md      |  2 +-
 .../zh_CN/source/tutorial/209-prompt_opt.md   |  2 +-
 .../conversation_with_langchain.py            |  4 ++--
 .../conversation_with_swe-agent/swe_agent.py  |  4 ++--
 .../distributed_debate/user_proxy_agent.py    |  5 +++--
 .../answerer_agent.py                         |  3 ++-
 .../searcher_agent.py                         |  4 +++-
 .../distributed_simulation/participant.py     |  8 +++++---
 examples/game_gomoku/code/board_agent.py      | 13 ++++++------
 examples/game_gomoku/code/gomoku_agent.py     |  4 ++--
 examples/game_gomoku/main.ipynb               |  7 ++++---
 src/agentscope/agents/agent.py                |  9 +++++----
 src/agentscope/agents/dialog_agent.py         | 16 +++++++--------
 src/agentscope/agents/dict_dialog_agent.py    | 17 ++++++++--------
 src/agentscope/agents/rag_agent.py            | 17 ++++++++--------
 src/agentscope/agents/react_agent.py          |  4 ++--
 src/agentscope/agents/rpc_agent.py            |  3 ++-
 src/agentscope/agents/text_to_image_agent.py  | 20 ++++++++++++++-----
 src/agentscope/agents/user_agent.py           | 16 +++++++--------
 src/agentscope/prompt/_prompt_comparer.py     |  4 ++--
 tests/msghub_test.py                          |  4 ++--
 tests/rpc_agent_test.py                       | 18 +++++++++--------
 30 files changed, 129 insertions(+), 99 deletions(-)

diff --git a/docs/sphinx_doc/en/source/tutorial/104-usecase.md b/docs/sphinx_doc/en/source/tutorial/104-usecase.md
index 94ba8ae92..7b0fddae1 100644
--- a/docs/sphinx_doc/en/source/tutorial/104-usecase.md
+++ b/docs/sphinx_doc/en/source/tutorial/104-usecase.md
@@ -46,9 +46,13 @@ To implement your own agent, you need to inherit `AgentBase` and implement the `
 
 ```python
 from agentscope.agents import AgentBase
+from agentscope.message import Msg
+
+
+from typing import Optional, Union, Sequence
 
 class MyAgent(AgentBase):
-    def reply(self, x):
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         # Do something here
         ...
         return x
diff --git a/docs/sphinx_doc/en/source/tutorial/201-agent.md b/docs/sphinx_doc/en/source/tutorial/201-agent.md
index 5b8563f4b..3fa916a88 100644
--- a/docs/sphinx_doc/en/source/tutorial/201-agent.md
+++ b/docs/sphinx_doc/en/source/tutorial/201-agent.md
@@ -39,7 +39,7 @@ class AgentBase(Operator):
     ) -> None:
 
     # ... [code omitted for brevity]
-    def observe(self, x: Union[dict, Sequence[dict]]) -> None:
+    def observe(self, x: Union[Msg, Sequence[Msg]]) -> None:
         # An optional method for updating the agent's internal state based on
         # messages it has observed. This method can be used to enrich the
         # agent's understanding and memory without producing an immediate
@@ -47,7 +47,7 @@ class AgentBase(Operator):
         if self.memory:
             self.memory.add(x)
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         # The core method to be implemented by custom agents. It defines the
         # logic for processing an input message and generating a suitable
         # response.
@@ -86,7 +86,7 @@ Below, we provide usages of how to configure various agents from the AgentPool:
 * **Reply Method**: The `reply` method is where the main logic for processing input *message* and generating responses.
 
 ```python
-def reply(self, x: dict = None) -> dict:
+def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
     # Additional processing steps can occur here
 
     # Record the input if needed
@@ -142,9 +142,9 @@ service_bot = DialogAgent(**dialog_agent_config)
 ```python
 def reply(
     self,
-    x: dict = None,
+    x: Optional[Union[Msg, Sequence[Msg]]] = None,
     required_keys: Optional[Union[list[str], str]] = None,
-) -> dict:
+) -> Msg:
     # Check if there is initial data to be added to memory
     if self.memory:
         self.memory.add(x)
diff --git a/docs/sphinx_doc/en/source/tutorial/203-parser.md b/docs/sphinx_doc/en/source/tutorial/203-parser.md
index 71fff4d59..bb2fae98e 100644
--- a/docs/sphinx_doc/en/source/tutorial/203-parser.md
+++ b/docs/sphinx_doc/en/source/tutorial/203-parser.md
@@ -211,7 +211,7 @@ In AgentScope, we achieve post-processing by calling the `to_content`, `to_memor
 
   ```python
       # ...
-      def reply(x: dict = None) -> None:
+      def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
 
           # ...
           res = self.model(prompt, parse_func=self.parser.parse)
diff --git a/docs/sphinx_doc/en/source/tutorial/204-service.md b/docs/sphinx_doc/en/source/tutorial/204-service.md
index dad6fa3d9..49b8b3b16 100644
--- a/docs/sphinx_doc/en/source/tutorial/204-service.md
+++ b/docs/sphinx_doc/en/source/tutorial/204-service.md
@@ -262,6 +262,9 @@ import json
 import inspect
 from agentscope.service import ServiceResponse
 from agentscope.agents import AgentBase
+from agentscope.message import Msg
+
+from typing import Optional, Union, Sequence
 
 
 def create_file(file_path: str, content: str = "") -> ServiceResponse:
@@ -282,7 +285,7 @@ def create_file(file_path: str, content: str = "") -> ServiceResponse:
 class YourAgent(AgentBase):
     # ... [omitted for brevity]
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         # ... [omitted for brevity]
 
         # construct a prompt to ask the agent to provide the parameters in JSON format
diff --git a/docs/sphinx_doc/en/source/tutorial/209-prompt_opt.md b/docs/sphinx_doc/en/source/tutorial/209-prompt_opt.md
index 9346ba5db..f1db3a248 100644
--- a/docs/sphinx_doc/en/source/tutorial/209-prompt_opt.md
+++ b/docs/sphinx_doc/en/source/tutorial/209-prompt_opt.md
@@ -397,6 +397,8 @@ from agentscope.agents import AgentBase
 from agentscope.prompt import SystemPromptOptimizer
 from agentscope.message import Msg
 
+from typing import Optional, Union, Sequence
+
 class MyAgent(AgentBase):
     def __init__(
             self,
@@ -411,7 +413,7 @@ class MyAgent(AgentBase):
             # or model_or_model_config_name=self.model
         )
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         self.memory.add(x)
 
         prompt = self.model.format(
diff --git a/docs/sphinx_doc/zh_CN/source/tutorial/104-usecase.md b/docs/sphinx_doc/zh_CN/source/tutorial/104-usecase.md
index 7727d0b0e..25608c845 100644
--- a/docs/sphinx_doc/zh_CN/source/tutorial/104-usecase.md
+++ b/docs/sphinx_doc/zh_CN/source/tutorial/104-usecase.md
@@ -47,9 +47,14 @@
 
 ```python
 from agentscope.agents import AgentBase
+from agentscope.message import Msg
+
+from typing import Optional, Union, Sequence
+
 
 class MyAgent(AgentBase):
-    def reply(self, x):
+
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         # Do something here
         ...
         return x
diff --git a/docs/sphinx_doc/zh_CN/source/tutorial/201-agent.md b/docs/sphinx_doc/zh_CN/source/tutorial/201-agent.md
index b6df3dc7a..10b29aeba 100644
--- a/docs/sphinx_doc/zh_CN/source/tutorial/201-agent.md
+++ b/docs/sphinx_doc/zh_CN/source/tutorial/201-agent.md
@@ -48,7 +48,7 @@ class AgentBase(Operator):
         if self.memory:
             self.memory.add(x)
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         # The core method to be implemented by custom agents. It defines the
         # logic for processing an input message and generating a suitable
         # response.
@@ -87,7 +87,7 @@ class AgentBase(Operator):
 * **回复方法**：`reply` 方法是处理输入消息和生成响应的主要逻辑所在
 
 ```python
-def reply(self, x: dict = None) -> dict:
+def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
     # Additional processing steps can occur here
 
     # Record the input if needed
@@ -143,9 +143,9 @@ service_bot = DialogAgent(**dialog_agent_config)
 ```python
 def reply(
     self,
-    x: dict = None,
+    x: Optional[Union[Msg, Sequence[Msg]]] = None,
     required_keys: Optional[Union[list[str], str]] = None,
-) -> dict:
+) -> Msg:
     # Check if there is initial data to be added to memory
     if self.memory:
         self.memory.add(x)
diff --git a/docs/sphinx_doc/zh_CN/source/tutorial/203-parser.md b/docs/sphinx_doc/zh_CN/source/tutorial/203-parser.md
index 1ae7fe863..8bad224e4 100644
--- a/docs/sphinx_doc/zh_CN/source/tutorial/203-parser.md
+++ b/docs/sphinx_doc/zh_CN/source/tutorial/203-parser.md
@@ -209,7 +209,7 @@ AgentScope中，我们通过调用`to_content`，`to_memory`和`to_metadata`方
 
   ```python
       # ...
-      def reply(x: dict = None) -> None:
+      def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
 
           # ...
           res = self.model(prompt, parse_func=self.parser.parse)
diff --git a/docs/sphinx_doc/zh_CN/source/tutorial/204-service.md b/docs/sphinx_doc/zh_CN/source/tutorial/204-service.md
index 788d2bdad..187612bc2 100644
--- a/docs/sphinx_doc/zh_CN/source/tutorial/204-service.md
+++ b/docs/sphinx_doc/zh_CN/source/tutorial/204-service.md
@@ -262,7 +262,7 @@ def create_file(file_path: str, content: str = "") -> ServiceResponse:
 class YourAgent(AgentBase):
     # ... [为简洁起见省略代码]
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         # ... [为简洁起见省略代码]
 
         # 构造提示，让代理提供 JSON 格式的参数
diff --git a/docs/sphinx_doc/zh_CN/source/tutorial/209-prompt_opt.md b/docs/sphinx_doc/zh_CN/source/tutorial/209-prompt_opt.md
index 47a972da5..7b7cd72a1 100644
--- a/docs/sphinx_doc/zh_CN/source/tutorial/209-prompt_opt.md
+++ b/docs/sphinx_doc/zh_CN/source/tutorial/209-prompt_opt.md
@@ -392,7 +392,7 @@ class MyAgent(AgentBase):
             # 或是 model_or_model_config_name=self.model
         )
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         self.memory.add(x)
 
         prompt = self.model.format(
diff --git a/examples/conversation_with_langchain/conversation_with_langchain.py b/examples/conversation_with_langchain/conversation_with_langchain.py
index 5d9f9c6b3..16b678892 100644
--- a/examples/conversation_with_langchain/conversation_with_langchain.py
+++ b/examples/conversation_with_langchain/conversation_with_langchain.py
@@ -2,7 +2,7 @@
 """A simple example of using langchain to create an assistant agent in
 AgentScope."""
 import os
-from typing import Optional
+from typing import Optional, Union, Sequence
 
 from langchain_openai import OpenAI
 from langchain.memory import ConversationBufferMemory
@@ -52,7 +52,7 @@ def __init__(self, name: str) -> None:
         )
         # [END] BY LANGCHAIN
 
-    def reply(self, x: Optional[dict] = None) -> Msg:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         # [START] BY LANGCHAIN
 
         # Generate response
diff --git a/examples/conversation_with_swe-agent/swe_agent.py b/examples/conversation_with_swe-agent/swe_agent.py
index 3b55431d5..6d2c49424 100644
--- a/examples/conversation_with_swe-agent/swe_agent.py
+++ b/examples/conversation_with_swe-agent/swe_agent.py
@@ -11,7 +11,7 @@
 from agentscope.message import Msg
 from agentscope.exception import ResponseParsingError
 from agentscope.parsers import MarkdownJsonDictParser
-from typing import List, Callable
+from typing import List, Callable, Optional, Union, Sequence
 import json
 from agentscope.service import (
     ServiceFactory,
@@ -206,7 +206,7 @@ def step(self) -> Msg:
         self.running_memory.append(str(action) + str(obs))
         return msg_res
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         action_name = None
         self.main_goal = x.content
         while not action_name == "exit":
diff --git a/examples/distributed_debate/user_proxy_agent.py b/examples/distributed_debate/user_proxy_agent.py
index cb27f3386..37f4b6d28 100644
--- a/examples/distributed_debate/user_proxy_agent.py
+++ b/examples/distributed_debate/user_proxy_agent.py
@@ -4,6 +4,7 @@
 from typing import Optional
 
 from agentscope.agents import UserAgent
+from agentscope.message import Msg
 
 
 class UserProxyAgent(UserAgent):
@@ -11,9 +12,9 @@ class UserProxyAgent(UserAgent):
 
     def reply(  # type: ignore[override]
         self,
-        x: dict = None,
+        x: Optional[Union[Msg, Sequence[Msg]]] = None,
         required_keys: Optional[Union[list[str], str]] = None,
-    ) -> dict:
+    ) -> Msg:
         """
         Reply with `self.speak(x)`
         """
diff --git a/examples/distributed_parallel_optimization/answerer_agent.py b/examples/distributed_parallel_optimization/answerer_agent.py
index 56ba014e9..e44551d01 100644
--- a/examples/distributed_parallel_optimization/answerer_agent.py
+++ b/examples/distributed_parallel_optimization/answerer_agent.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 """Answerer Agent."""
+from typing import Optional, Union, Sequence
 
 from agentscope.message import Msg
 from agentscope.agents import AgentBase
@@ -22,7 +23,7 @@ def __init__(
             use_memory=False,
         )
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         response = load_web(
             url=x.url,
             keep_raw=False,
diff --git a/examples/distributed_parallel_optimization/searcher_agent.py b/examples/distributed_parallel_optimization/searcher_agent.py
index 9a1c9b435..eb1ad2f23 100644
--- a/examples/distributed_parallel_optimization/searcher_agent.py
+++ b/examples/distributed_parallel_optimization/searcher_agent.py
@@ -2,6 +2,8 @@
 """Searcher agent."""
 
 from functools import partial
+from typing import Optional, Union, Sequence
+
 from agentscope.message import Msg
 from agentscope.agents import AgentBase
 from agentscope.service import google_search, bing_search
@@ -56,7 +58,7 @@ def __init__(
             assert api_key is not None, "bing search requires 'api_key'"
             self.search = partial(bing_search, api_key=api_key)
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         prompt = self.model.format(
             Msg(name="system", role="system", content=self.sys_prompt),
             x,
diff --git a/examples/distributed_simulation/participant.py b/examples/distributed_simulation/participant.py
index dac3d17bf..f017d4de3 100644
--- a/examples/distributed_simulation/participant.py
+++ b/examples/distributed_simulation/participant.py
@@ -3,6 +3,8 @@
 import random
 import time
 import re
+from typing import Optional, Union, Sequence
+
 from loguru import logger
 
 from agentscope.message import Msg
@@ -30,7 +32,7 @@ def generate_random_response(self) -> str:
         time.sleep(self.sleep_time)
         return str(random.randint(0, self.max_value))
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """Generate a random value"""
         # generate a response in content
         response = self.generate_random_response()
@@ -74,7 +76,7 @@ def parse_value(self, txt: str) -> str:
         else:
             return numbers[-1]
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """Generate a value by LLM"""
         if self.memory:
             self.memory.add(x)
@@ -134,7 +136,7 @@ def __init__(
                 for config in part_configs
             ]
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         results = []
         msg = Msg(
             name="moderator",
diff --git a/examples/game_gomoku/code/board_agent.py b/examples/game_gomoku/code/board_agent.py
index 6cbef4ced..b0cde430f 100644
--- a/examples/game_gomoku/code/board_agent.py
+++ b/examples/game_gomoku/code/board_agent.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 """A board agent class that can host a Gomoku game, and a function to
 convert the board to an image."""
+from typing import Optional, Union, Sequence
 
 import numpy as np
 from matplotlib import pyplot as plt, patches
@@ -81,7 +82,7 @@ def __init__(self, name: str) -> None:
         # Record the status of the game
         self.game_end = False
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         if x is None:
             # Beginning of the game
             content = (
@@ -89,12 +90,12 @@ def reply(self, x: dict = None) -> dict:
                 "Please make your move."
             )
         else:
-            row, col = x["content"]
+            row, col = x.content
 
             self.assert_valid_move(row, col)
 
             # change the board
-            self.board[row, col] = NAME_TO_PIECE[x["name"]]
+            self.board[row, col] = NAME_TO_PIECE[x.name]
 
             # check if the game ends
             if self.check_draw():
@@ -102,15 +103,15 @@ def reply(self, x: dict = None) -> dict:
                 self.game_end = True
             else:
                 next_player_name = (
-                    NAME_BLACK if x["name"] == NAME_WHITE else NAME_WHITE
+                    NAME_BLACK if x.name == NAME_WHITE else NAME_WHITE
                 )
                 content = CURRENT_BOARD_PROMPT_TEMPLATE.format(
                     board=self.board2text(),
                     player=next_player_name,
                 )
 
-                if self.check_win(row, col, NAME_TO_PIECE[x["name"]]):
-                    content = f"The game ends, {x['name']} wins!"
+                if self.check_win(row, col, NAME_TO_PIECE[x.name]):
+                    content = f"The game ends, {x.name} wins!"
                     self.game_end = True
 
         msg_host = Msg(self.name, content, role="assistant")
diff --git a/examples/game_gomoku/code/gomoku_agent.py b/examples/game_gomoku/code/gomoku_agent.py
index 16e05cc36..0615f7909 100644
--- a/examples/game_gomoku/code/gomoku_agent.py
+++ b/examples/game_gomoku/code/gomoku_agent.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 """A Gomoku agent that can play the game with another agent."""
 
-from typing import Optional
+from typing import Optional, Union, Sequence
 
 import json
 
@@ -48,7 +48,7 @@ def __init__(
 
         self.memory.add(Msg("system", sys_prompt, role="system"))
 
-    def reply(self, x: Optional[dict] = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         if self.memory:
             self.memory.add(x)
 
diff --git a/examples/game_gomoku/main.ipynb b/examples/game_gomoku/main.ipynb
index b3b31831a..a0e3d996f 100644
--- a/examples/game_gomoku/main.ipynb
+++ b/examples/game_gomoku/main.ipynb
@@ -130,6 +130,7 @@
     "import numpy as np\n",
     "from agentscope.message import Msg\n",
     "from agentscope.agents import AgentBase\n",
+    "from typing import Optional, Union, Sequence\n",
     "\n",
     "CURRENT_BOARD_PROMPT_TEMPLATE = \"\"\"The current board is as follows:\n",
     "{board}\n",
@@ -158,7 +159,7 @@
     "        # Record the status of the game\n",
     "        self.game_end = False\n",
     "        \n",
-    "    def reply(self, x: dict = None) -> dict:\n",
+    "    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:\n",
     "        if x is None:\n",
     "            # Beginning of the game\n",
     "            content = (\n",
@@ -263,7 +264,7 @@
    "source": [
     "import json\n",
     "from agentscope.models import ModelResponse\n",
-    "from typing import Optional\n",
+    "from typing import Optional, Union, Sequence\n",
     "\n",
     "SYS_PROMPT_TEMPLATE = \"\"\"\n",
     "You're a skillful Gomoku player. You should play against your opponent according to the following rules:\n",
@@ -304,7 +305,7 @@
     "        \n",
     "        self.memory.add(Msg(\"system\", sys_prompt, role=\"system\"))\n",
     "        \n",
-    "    def reply(self, x: Optional[dict] = None) -> dict:\n",
+    "    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:\n",
     "        if self.memory:\n",
     "            self.memory.add(x)\n",
     "        \n",
diff --git a/src/agentscope/agents/agent.py b/src/agentscope/agents/agent.py
index 7ae48a5b7..a7a25485e 100644
--- a/src/agentscope/agents/agent.py
+++ b/src/agentscope/agents/agent.py
@@ -252,15 +252,16 @@ def register_agent_class(cls, agent_class: Type[AgentBase]) -> None:
         else:
             cls._registry[agent_class_name] = agent_class
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """Define the actions taken by this agent.
 
         Args:
-            x (`dict`, defaults to `None`):
-                Dialog history and some environment information
+            x (`Optional[Union[Msg, Sequence[Msg]]]`, defaults to `None`):
+                The input message(s) to the agent, which also can be omitted if
+                the agent doesn't need any input.
 
         Returns:
-            The agent's response to the input.
+            `Msg`: The output message generated by the agent.
 
         Note:
             Given that some agents are in an adversarial environment,
diff --git a/src/agentscope/agents/dialog_agent.py b/src/agentscope/agents/dialog_agent.py
index 1ece870b6..28f50d58a 100644
--- a/src/agentscope/agents/dialog_agent.py
+++ b/src/agentscope/agents/dialog_agent.py
@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 """A general dialog agent."""
-from typing import Optional
+from typing import Optional, Union, Sequence
 
 from ..message import Msg
 from .agent import AgentBase
@@ -42,21 +42,19 @@ def __init__(
             memory_config=memory_config,
         )
 
-    # TODO change typing from dict to MSG
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """Reply function of the agent. Processes the input data,
         generates a prompt using the current dialogue memory and system
         prompt, and invokes the language model to produce a response. The
         response is then formatted and added to the dialogue memory.
 
         Args:
-            x (`dict`, defaults to `None`):
-                A dictionary representing the user's input to the agent. This
-                input is added to the dialogue memory if provided. Defaults to
-                None.
+            x (`Optional[Union[Msg, Sequence[Msg]]]`, defaults to `None`):
+                The input message(s) to the agent, which also can be omitted if
+                the agent doesn't need any input.
+
         Returns:
-            A dictionary representing the message generated by the agent in
-            response to the user's input.
+            `Msg`: The output message generated by the agent.
         """
         # record the input if needed
         if self.memory:
diff --git a/src/agentscope/agents/dict_dialog_agent.py b/src/agentscope/agents/dict_dialog_agent.py
index eb16690e0..5ce4ce065 100644
--- a/src/agentscope/agents/dict_dialog_agent.py
+++ b/src/agentscope/agents/dict_dialog_agent.py
@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 """An agent that replies in a dictionary format."""
-from typing import Optional
+from typing import Optional, Union, Sequence
 
 from ..message import Msg
 from .agent import AgentBase
@@ -64,7 +64,7 @@ def set_parser(self, parser: ParserBase) -> None:
         """
         self.parser = parser
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """Reply function of the agent.
         Processes the input data, generates a prompt using the current
         dialogue memory and system prompt, and invokes the language
@@ -72,14 +72,13 @@ def reply(self, x: dict = None) -> dict:
         and added to the dialogue memory.
 
         Args:
-            x (`dict`, defaults to `None`):
-                A dictionary representing the user's input to the agent.
-                This input is added to the dialogue memory if provided.
+            x (`Optional[Union[Msg, Sequence[Msg]]]`, defaults to `None`):
+                The input message(s) to the agent, which also can be omitted if
+                the agent doesn't need any input.
+
+
         Returns:
-            A dictionary representing the message generated by the agent in
-            response to the user's input. It contains at least a 'speak' key
-            with the textual response and may include other keys such as
-            'agreement' if provided by the language model.
+            `Msg`: The output message generated by the agent.
 
         Raises:
             `json.decoder.JSONDecodeError`:
diff --git a/src/agentscope/agents/rag_agent.py b/src/agentscope/agents/rag_agent.py
index 59012b17c..63a23fdcd 100644
--- a/src/agentscope/agents/rag_agent.py
+++ b/src/agentscope/agents/rag_agent.py
@@ -6,7 +6,7 @@
 Notice, this is a Beta version of RAG agent.
 """
 
-from typing import Any
+from typing import Any, Optional, Union, Sequence
 from loguru import logger
 
 from agentscope.agents.agent import AgentBase
@@ -82,7 +82,7 @@ def __init__(
         self.recent_n_mem_for_retrieve = recent_n_mem_for_retrieve
         self.description = kwargs.get("description", "")
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """
         Reply function of the RAG agent.
         Processes the input data,
@@ -93,13 +93,12 @@ def reply(self, x: dict = None) -> dict:
         response is then formatted and added to the dialogue memory.
 
         Args:
-            x (`dict`, defaults to `None`):
-                A dictionary representing the user's input to the agent. This
-                input is added to the memory if provided. Defaults to
-                None.
+            x (`Optional[Union[Msg, Sequence[Msg]]]`, defaults to `None`):
+                The input message(s) to the agent, which also can be omitted if
+                the agent doesn't need any input.
+
         Returns:
-            A dictionary representing the message generated by the agent in
-            response to the user's input.
+            `Msg`: The output message generated by the agent.
         """
         retrieved_docs_to_string = ""
         # record the input if needed
@@ -118,7 +117,7 @@ def reply(self, x: dict = None) -> dict:
                 else str(history)
             )
         elif x is not None:
-            query = x["content"]
+            query = x.content
         else:
             query = ""
 
diff --git a/src/agentscope/agents/react_agent.py b/src/agentscope/agents/react_agent.py
index cdc81788b..cf0522e70 100644
--- a/src/agentscope/agents/react_agent.py
+++ b/src/agentscope/agents/react_agent.py
@@ -3,7 +3,7 @@
 and act iteratively to solve problems. More details can be found in the paper
 https://arxiv.org/abs/2210.03629.
 """
-from typing import Any
+from typing import Any, Optional, Union, Sequence
 
 from loguru import logger
 
@@ -140,7 +140,7 @@ def __init__(
             keys_to_content=True if self.verbose else "speak",
         )
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """The reply function that achieves the ReAct algorithm.
         The more details please refer to https://arxiv.org/abs/2210.03629"""
 
diff --git a/src/agentscope/agents/rpc_agent.py b/src/agentscope/agents/rpc_agent.py
index ec6cc39ba..b175fa781 100644
--- a/src/agentscope/agents/rpc_agent.py
+++ b/src/agentscope/agents/rpc_agent.py
@@ -6,6 +6,7 @@
 from agentscope.message import (
     PlaceholderMessage,
     serialize,
+    Msg,
 )
 from agentscope.rpc import RpcAgentClient
 from agentscope.server.launcher import RpcAgentServerLauncher
@@ -105,7 +106,7 @@ def _launch_server(self) -> None:
         )
         self.client.create_agent(self.agent_configs)
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         if self.client is None:
             self._launch_server()
         return PlaceholderMessage(
diff --git a/src/agentscope/agents/text_to_image_agent.py b/src/agentscope/agents/text_to_image_agent.py
index f31fb55df..00519a404 100644
--- a/src/agentscope/agents/text_to_image_agent.py
+++ b/src/agentscope/agents/text_to_image_agent.py
@@ -1,7 +1,9 @@
 # -*- coding: utf-8 -*-
 """An agent that convert text to image."""
 
-from typing import Optional
+from typing import Optional, Union, Sequence
+
+from loguru import logger
 
 from .agent import AgentBase
 from ..message import Msg
@@ -42,7 +44,12 @@ def __init__(
             memory_config=memory_config,
         )
 
-    def reply(self, x: dict = None) -> dict:
+        logger.warning(
+            "The `TextToImageAgent` will be deprecated in v0.0.6, "
+            "please use `text_to_image` service and `ReActAgent` instead.",
+        )
+
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         if self.memory:
             self.memory.add(x)
         if x is None:
@@ -50,13 +57,16 @@ def reply(self, x: dict = None) -> dict:
             if self.memory and self.memory.size() > 0:
                 x = self.memory.get_memory()[-1]
             else:
-                # if no message find, just return None
-                return {}
+                return Msg(
+                    self.name,
+                    content="Please provide a text prompt to generate image.",
+                    role="assistant",
+                )
         image_urls = self.model(x.content).image_urls
         # TODO: optimize the construction of content
         msg = Msg(
             self.name,
-            content="This is the generated image ",
+            content="This is the generated image",
             role="assistant",
             url=image_urls,
         )
diff --git a/src/agentscope/agents/user_agent.py b/src/agentscope/agents/user_agent.py
index 3d0b1cf3e..b76cf28d5 100644
--- a/src/agentscope/agents/user_agent.py
+++ b/src/agentscope/agents/user_agent.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 """User Agent class"""
 import time
-from typing import Union
+from typing import Union, Sequence
 from typing import Optional
 from loguru import logger
 
@@ -33,10 +33,10 @@ def __init__(self, name: str = "User", require_url: bool = False) -> None:
 
     def reply(
         self,
-        x: dict = None,
+        x: Optional[Union[Msg, Sequence[Msg]]] = None,
         required_keys: Optional[Union[list[str], str]] = None,
         timeout: Optional[int] = None,
-    ) -> dict:
+    ) -> Msg:
         """
         Processes the input provided by the user and stores it in memory,
         potentially formatting it with additional provided details.
@@ -47,9 +47,9 @@ def reply(
         added to the object's memory.
 
         Arguments:
-            x (`dict`, defaults to `None`):
-                A dictionary containing initial data to be added to memory.
-                Defaults to None.
+            x (`Optional[Union[Msg, Sequence[Msg]]]`, defaults to `None`):
+                The input message(s) to the agent, which also can be omitted if
+                the agent doesn't need any input.
             required_keys \
                 (`Optional[Union[list[str], str]]`, defaults to `None`):
                 Strings that requires user to input, which will be used as
@@ -59,9 +59,7 @@ def reply(
                 for no limit.
 
         Returns:
-            `dict`: A dictionary representing the message object that contains
-            the user's input and any additional details. This is also
-            stored in the object's memory.
+            `Msg`: The output message generated by the agent.
         """
         if self.memory:
             self.memory.add(x)
diff --git a/src/agentscope/prompt/_prompt_comparer.py b/src/agentscope/prompt/_prompt_comparer.py
index c1d6c4dce..bc382eae2 100644
--- a/src/agentscope/prompt/_prompt_comparer.py
+++ b/src/agentscope/prompt/_prompt_comparer.py
@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 """The Abtest module to show how different system prompt performs"""
-from typing import List
+from typing import List, Optional, Union, Sequence
 from loguru import logger
 
 from agentscope.models import load_model_by_config_name
@@ -41,7 +41,7 @@ def enable_display(self) -> None:
         """Enable the display of the output message."""
         self.display = True
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """Reply the message with the given system prompt."""
         self.memory.add(x)
 
diff --git a/tests/msghub_test.py b/tests/msghub_test.py
index 75f61dccc..9859c364e 100644
--- a/tests/msghub_test.py
+++ b/tests/msghub_test.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 """ Unit test for msghub."""
 import unittest
-from typing import Optional
+from typing import Optional, Union, Sequence
 
 from agentscope.agents import AgentBase
 from agentscope import msghub
@@ -11,7 +11,7 @@
 class TestAgent(AgentBase):
     """Test agent class for msghub."""
 
-    def reply(self, x: Optional[dict] = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """Reply function for agent."""
         if x is not None:
             self.memory.add(x)
diff --git a/tests/rpc_agent_test.py b/tests/rpc_agent_test.py
index 321087a44..7462dd388 100644
--- a/tests/rpc_agent_test.py
+++ b/tests/rpc_agent_test.py
@@ -6,6 +6,8 @@
 import time
 import os
 import shutil
+from typing import Optional, Union, Sequence
+
 from loguru import logger
 
 import agentscope
@@ -26,7 +28,7 @@ def __init__(self, **kwargs) -> None:  # type: ignore[no-untyped-def]
         super().__init__(**kwargs)
         self.id = 0
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """Response after 2s"""
         x.id = self.id
         self.id += 1
@@ -37,7 +39,7 @@ def reply(self, x: dict = None) -> dict:
 class DemoRpcAgentAdd(AgentBase):
     """A demo Rpc agent for test usage"""
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """add the value, wait 1s"""
         x.content["value"] += 1
         time.sleep(1)
@@ -47,7 +49,7 @@ def reply(self, x: dict = None) -> dict:
 class DemoLocalAgentAdd(AgentBase):
     """A demo local agent for test usage"""
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         """add the value, wait 1s"""
         x.content["value"] += 1
         time.sleep(1)
@@ -57,7 +59,7 @@ def reply(self, x: dict = None) -> dict:
 class DemoRpcAgentWithMemory(AgentBase):
     """A demo Rpc agent that count its memory"""
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         self.memory.add(x)
         msg = Msg(
             name=self.name,
@@ -72,7 +74,7 @@ def reply(self, x: dict = None) -> dict:
 class DemoRpcAgentWithMonitor(AgentBase):
     """A demo Rpc agent that use monitor"""
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         monitor = MonitorFactory.get_monitor()
         try:
             monitor.update({"msg_num": 1})
@@ -103,7 +105,7 @@ def __init__(self, name: str, value: int) -> None:
         super().__init__(name)
         self.value = value
 
-    def reply(self, _: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         time.sleep(1)
         return Msg(
             name=self.name,
@@ -126,7 +128,7 @@ def __init__(
         super().__init__(name, to_dist=to_dist)
         self.agents = agents
 
-    def reply(self, _: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         result = []
         stime = time.time()
         for agent in self.agents:
@@ -148,7 +150,7 @@ def reply(self, _: dict = None) -> dict:
 class DemoErrorAgent(AgentBase):
     """A demo Rpc agent that raise Error"""
 
-    def reply(self, x: dict = None) -> dict:
+    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None) -> Msg:
         raise RuntimeError("Demo Error")