change message replication

InternLM · Oct 29, 2024 · 97d4020 · 97d4020
1 parent cd34d8d
commit 97d4020
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 16 deletions.
diff --git a/README.md b/README.md
@@ -367,12 +367,9 @@ class PrefixedMessageHook(Hook):
         self.senders = senders or []
 
     def before_agent(self, agent, messages, session_id):
-        for i, message in enumerate(messages):
+        for message in messages:
             if message.sender in self.senders:
-                message = message.copy(deep=True)
                 message.content = self.prefix + message.content
-                messages[i] = message
-        return messages
 
 class AsyncBlogger(AsyncAgent):
     def __init__(self, model_path, writer_prompt, critic_prompt, critic_prefix='', max_turn=3):

diff --git a/docs/en/get_started/quickstart.md b/docs/en/get_started/quickstart.md
@@ -339,12 +339,9 @@ class PrefixedMessageHook(Hook):
         self.senders = senders or []
 
     def before_agent(self, agent, messages, session_id):
-        for i, message in enumerate(messages):
+        for message in messages:
             if message.sender in self.senders:
-                message = message.copy(deep=True)
                 message.content = self.prefix + message.content
-                messages[i] = message
-        return messages
 
 class AsyncBlogger(AsyncAgent):
     def __init__(self, model_path, writer_prompt, critic_prompt, critic_prefix='', max_turn=3):

diff --git a/lagent/agents/agent.py b/lagent/agents/agent.py
@@ -70,11 +70,10 @@ def __call__(
     ) -> AgentMessage:
         # message.receiver = self.name
         message = [
-            AgentMessage(sender='user', content=m) if isinstance(m, str) else m
-            for m in message
+            AgentMessage(sender='user', content=m)
+            if isinstance(m, str) else copy.deepcopy(m) for m in message
         ]
         for hook in self._hooks.values():
-            message = copy.deepcopy(message)
             result = hook.before_agent(self, message, session_id)
             if result:
                 message = result
@@ -87,8 +86,8 @@ def __call__(
                 content=response_message,
             )
         self.update_memory(response_message, session_id=session_id)
+        response_message = copy.deepcopy(response_message)
         for hook in self._hooks.values():
-            response_message = copy.deepcopy(response_message)
             result = hook.after_agent(self, response_message, session_id)
             if result:
                 response_message = result
@@ -177,11 +176,10 @@ async def __call__(self,
                        session_id=0,
                        **kwargs) -> AgentMessage:
         message = [
-            AgentMessage(sender='user', content=m) if isinstance(m, str) else m
-            for m in message
+            AgentMessage(sender='user', content=m)
+            if isinstance(m, str) else copy.deepcopy(m) for m in message
         ]
         for hook in self._hooks.values():
-            message = copy.deepcopy(message)
             result = hook.before_agent(self, message, session_id)
             if result:
                 message = result
@@ -194,8 +192,8 @@ async def __call__(self,
                 content=response_message,
             )
         self.update_memory(response_message, session_id=session_id)
+        response_message = copy.deepcopy(response_message)
         for hook in self._hooks.values():
-            response_message = copy.deepcopy(response_message)
             result = hook.after_agent(self, response_message, session_id)
             if result:
                 response_message = result