Enhancing RolePlaying with custom ChatAgent support (#3199)

Co-authored-by: Adam Chen <[email protected]>
Co-authored-by: Saed Bhati <[email protected]>
Co-authored-by: Waleed Alzarooni <[email protected]>
Co-authored-by: JINO ROHIT <[email protected]>

Author: 5 people

camel/societies/role_playing.py

@@ -62,8 +62,8 @@ class RolePlaying:
             (default: :obj:`TaskType.AI_SOCIETY`)
         assistant_agent_kwargs (Dict, optional): Additional arguments to pass
             to the assistant agent. (default: :obj:`None`)
-        user_agent_kwargs (Dict, optional): Additional arguments to pass to
-            the user agent. (default: :obj:`None`)
+        user_agent_kwargs (Dict, optional): Additional arguments to pass to the
+            user agent. (default: :obj:`None`)
         task_specify_agent_kwargs (Dict, optional): Additional arguments to
             pass to the task specify agent. (default: :obj:`None`)
         task_planner_agent_kwargs (Dict, optional): Additional arguments to
@@ -81,6 +81,12 @@ class RolePlaying:
         stop_event (Optional[threading.Event], optional): Event to signal
             termination of the agent's operation. When set, the agent will
             terminate its execution. (default: :obj:`None`)
+        assistant_agent (ChatAgent, optional): A pre-configured ChatAgent to
+            use as the assistant. If provided, this will override the creation
+            of a new assistant agent. (default: :obj:`None`)
+        user_agent (ChatAgent, optional): A pre-configured ChatAgent to use as
+            the user. If provided, this will override the creation of a new
+            user agent. (default: :obj:`None`)
     """
 
     def __init__(
@@ -106,6 +112,8 @@ def __init__(
         extend_task_specify_meta_dict: Optional[Dict] = None,
         output_language: Optional[str] = None,
         stop_event: Optional[threading.Event] = None,
+        assistant_agent: Optional[ChatAgent] = None,
+        user_agent: Optional[ChatAgent] = None,
     ) -> None:
         if model is not None:
             logger.warning(
@@ -143,29 +151,56 @@ def __init__(
             **(sys_msg_generator_kwargs or {}),
         )
 
-        (
-            init_assistant_sys_msg,
-            init_user_sys_msg,
-            sys_msg_meta_dicts,
-        ) = self._get_sys_message_info(
-            assistant_role_name,
-            user_role_name,
-            sys_msg_generator,
-            extend_sys_msg_meta_dicts=extend_sys_msg_meta_dicts,
-        )
-
+        # Initialize agent attributes first
         self.assistant_agent: ChatAgent
         self.user_agent: ChatAgent
-        self.assistant_sys_msg: Optional[BaseMessage]
-        self.user_sys_msg: Optional[BaseMessage]
+        self.assistant_sys_msg: Optional[BaseMessage] = None
+        self.user_sys_msg: Optional[BaseMessage] = None
+
+        # Determine if we need to generate system messages
+        if assistant_agent is None or user_agent is None:
+            # Generate system messages for missing agents
+            (
+                init_assistant_sys_msg,
+                init_user_sys_msg,
+                sys_msg_meta_dicts,
+            ) = self._get_sys_message_info(
+                assistant_role_name,
+                user_role_name,
+                sys_msg_generator,
+                extend_sys_msg_meta_dicts=extend_sys_msg_meta_dicts,
+            )
+        else:
+            # When both agents are provided, use their existing system messages
+            assistant_sys_msg = assistant_agent.system_message
+            user_sys_msg = user_agent.system_message
+
+            # Ensure system messages are not None
+            if assistant_sys_msg is None:
+                raise ValueError(
+                    "Provided assistant_agent has None system_message"
+                )
+            if user_sys_msg is None:
+                raise ValueError("Provided user_agent has None system_message")
+
+            init_assistant_sys_msg = assistant_sys_msg
+            init_user_sys_msg = user_sys_msg
+            # Create a default sys_msg_meta_dicts for critic initialization
+            sys_msg_meta_dicts = [
+                dict(task=self.task_prompt) for _ in range(2)
+            ]
+
         self._init_agents(
             init_assistant_sys_msg,
             init_user_sys_msg,
             assistant_agent_kwargs=assistant_agent_kwargs,
             user_agent_kwargs=user_agent_kwargs,
             output_language=output_language,
             stop_event=stop_event,
+            assistant_agent=assistant_agent,
+            user_agent=user_agent,
         )
+
         self.critic: Optional[Union[CriticAgent, Human]] = None
         self.critic_sys_msg: Optional[BaseMessage] = None
         self._init_critic(
@@ -320,20 +355,22 @@ def _get_sys_message_info(
 
     def _init_agents(
         self,
-        init_assistant_sys_msg: BaseMessage,
-        init_user_sys_msg: BaseMessage,
+        init_assistant_sys_msg: Optional[BaseMessage],
+        init_user_sys_msg: Optional[BaseMessage],
         assistant_agent_kwargs: Optional[Dict] = None,
         user_agent_kwargs: Optional[Dict] = None,
         output_language: Optional[str] = None,
         stop_event: Optional[threading.Event] = None,
+        assistant_agent: Optional[ChatAgent] = None,
+        user_agent: Optional[ChatAgent] = None,
     ) -> None:
         r"""Initialize assistant and user agents with their system messages.
 
         Args:
-            init_assistant_sys_msg (BaseMessage): Assistant agent's initial
+            init_assistant_sys_msg (Optional[BaseMessage]): Assistant agent's
+                initial system message.
+            init_user_sys_msg (Optional[BaseMessage]): User agent's initial
                 system message.
-            init_user_sys_msg (BaseMessage): User agent's initial system
-                message.
             assistant_agent_kwargs (Dict, optional): Additional arguments to
                 pass to the assistant agent. (default: :obj:`None`)
             user_agent_kwargs (Dict, optional): Additional arguments to
@@ -343,6 +380,12 @@ def _init_agents(
             stop_event (Optional[threading.Event], optional): Event to signal
                 termination of the agent's operation. When set, the agent will
                 terminate its execution. (default: :obj:`None`)
+            assistant_agent (ChatAgent, optional): A pre-configured ChatAgent
+                to use as the assistant. If provided, this will override the
+                creation of a new assistant agent. (default: :obj:`None`)
+            user_agent (ChatAgent, optional): A pre-configured ChatAgent to use
+                as the user. If provided, this will override the creation of a
+                new user agent. (default: :obj:`None`)
         """
         if self.model is not None:
             if assistant_agent_kwargs is None:
@@ -354,21 +397,71 @@ def _init_agents(
             elif 'model' not in user_agent_kwargs:
                 user_agent_kwargs.update(dict(model=self.model))
 
-        self.assistant_agent = ChatAgent(
-            init_assistant_sys_msg,
-            output_language=output_language,
-            stop_event=stop_event,
-            **(assistant_agent_kwargs or {}),
-        )
-        self.assistant_sys_msg = self.assistant_agent.system_message
-
-        self.user_agent = ChatAgent(
-            init_user_sys_msg,
-            output_language=output_language,
-            stop_event=stop_event,
-            **(user_agent_kwargs or {}),
-        )
-        self.user_sys_msg = self.user_agent.system_message
+        # Use provided assistant agent if available, otherwise create a new one
+        if assistant_agent is not None:
+            # Ensure functionality consistent with our configuration
+            if (
+                hasattr(assistant_agent, 'output_language')
+                and output_language is not None
+            ):
+                assistant_agent.output_language = output_language
+            if hasattr(assistant_agent, 'stop_event'):
+                assistant_agent.stop_event = stop_event
+            self.assistant_agent = assistant_agent
+            # Handle potential None system_message - use provided or fallback
+            if assistant_agent.system_message is not None:
+                self.assistant_sys_msg = assistant_agent.system_message
+            elif init_assistant_sys_msg is not None:
+                self.assistant_sys_msg = init_assistant_sys_msg
+            else:
+                raise ValueError("Assistant system message cannot be None")
+        else:
+            # Create new assistant agent
+            if init_assistant_sys_msg is None:
+                raise ValueError(
+                    "Assistant system message cannot be None when creating "
+                    "new agent"
+                )
+            self.assistant_agent = ChatAgent(
+                init_assistant_sys_msg,
+                output_language=output_language,
+                stop_event=stop_event,
+                **(assistant_agent_kwargs or {}),
+            )
+            self.assistant_sys_msg = self.assistant_agent.system_message
+
+        # Use provided user agent if available, otherwise create a new one
+        if user_agent is not None:
+            # Ensure functionality consistent with our configuration
+            if (
+                hasattr(user_agent, 'output_language')
+                and output_language is not None
+            ):
+                user_agent.output_language = output_language
+            if hasattr(user_agent, 'stop_event'):
+                user_agent.stop_event = stop_event
+            self.user_agent = user_agent
+            # Handle potential None system_message - use provided or fallback
+            if user_agent.system_message is not None:
+                self.user_sys_msg = user_agent.system_message
+            elif init_user_sys_msg is not None:
+                self.user_sys_msg = init_user_sys_msg
+            else:
+                raise ValueError("User system message cannot be None")
+        else:
+            # Create new user agent
+            if init_user_sys_msg is None:
+                raise ValueError(
+                    "User system message cannot be None when creating new "
+                    "agent"
+                )
+            self.user_agent = ChatAgent(
+                init_user_sys_msg,
+                output_language=output_language,
+                stop_event=stop_event,
+                **(user_agent_kwargs or {}),
+            )
+            self.user_sys_msg = self.user_agent.system_message
 
     def _init_critic(
         self,
@@ -389,7 +482,7 @@ def _init_critic(
             sys_msg_meta_dicts (list): A list of system message meta dicts.
             critic_role_name (str): The name of the role played by the critic.
             critic_criteria (str, optional): Critic criteria for the
-                critic agent. If not specified, set the criteria to
+                critic agent. If not specified, set it to
                 improve task performance. (default: :obj:`None`)
             critic_kwargs (Dict, optional): Additional arguments to
                 pass to the critic. (default: :obj:`None`)
@@ -465,20 +558,28 @@ def init_chat(self, init_msg_content: Optional[str] = None) -> BaseMessage:
             BaseMessage: A single `BaseMessage` representing the initial
                 message.
         """
-        self.assistant_agent.reset()
-        self.user_agent.reset()
+        if self.assistant_agent is not None:
+            self.assistant_agent.reset()
+        if self.user_agent is not None:
+            self.user_agent.reset()
         default_init_msg_content = (
             "Now start to give me instructions one by one. "
             "Only reply with Instruction and Input."
         )
-        if init_msg_content is None:
-            init_msg_content = default_init_msg_content
+        final_init_msg_content = init_msg_content or default_init_msg_content
 
         # Initialize a message sent by the assistant
+        assistant_role_name = "assistant"
+        if self.assistant_sys_msg is not None and hasattr(
+            self.assistant_sys_msg, 'role_name'
+        ):
+            role_name_attr = getattr(self.assistant_sys_msg, 'role_name', None)
+            if role_name_attr is not None:
+                assistant_role_name = str(role_name_attr)
+
         init_msg = BaseMessage.make_assistant_message(
-            role_name=getattr(self.assistant_sys_msg, 'role_name', None)
-            or "assistant",
-            content=init_msg_content,
+            role_name=assistant_role_name,
+            content=final_init_msg_content,
         )
 
         return init_msg
@@ -501,20 +602,28 @@ async def ainit_chat(
         """
         # Currently, reset() is synchronous, but if it becomes async in the
         # future, we can await it here
-        self.assistant_agent.reset()
-        self.user_agent.reset()
+        if self.assistant_agent is not None:
+            self.assistant_agent.reset()
+        if self.user_agent is not None:
+            self.user_agent.reset()
         default_init_msg_content = (
             "Now start to give me instructions one by one. "
             "Only reply with Instruction and Input."
         )
-        if init_msg_content is None:
-            init_msg_content = default_init_msg_content
+        final_init_msg_content = init_msg_content or default_init_msg_content
 
         # Initialize a message sent by the assistant
+        assistant_role_name = "assistant"
+        if self.assistant_sys_msg is not None and hasattr(
+            self.assistant_sys_msg, 'role_name'
+        ):
+            role_name_attr = getattr(self.assistant_sys_msg, 'role_name', None)
+            if role_name_attr is not None:
+                assistant_role_name = str(role_name_attr)
+
         init_msg = BaseMessage.make_assistant_message(
-            role_name=getattr(self.assistant_sys_msg, 'role_name', None)
-            or "assistant",
-            content=init_msg_content,
+            role_name=assistant_role_name,
+            content=final_init_msg_content,
         )
 
         return init_msg
@@ -544,6 +653,11 @@ def step(
                 user agent terminated the conversation, and any additional user
                 information.
         """
+        if self.user_agent is None:
+            raise ValueError("User agent is not initialized")
+        if self.assistant_agent is None:
+            raise ValueError("Assistant agent is not initialized")
+
         user_response = self.user_agent.step(assistant_msg)
         if user_response.terminated or user_response.msgs is None:
             return (
@@ -620,6 +734,11 @@ async def astep(
                 user agent terminated the conversation, and any additional user
                 information.
         """
+        if self.user_agent is None:
+            raise ValueError("User agent is not initialized")
+        if self.assistant_agent is None:
+            raise ValueError("Assistant agent is not initialized")
+
         user_response = await self.user_agent.astep(assistant_msg)
         if user_response.terminated or user_response.msgs is None:
             return (
@@ -682,6 +801,10 @@ def clone(
             RolePlaying: A new instance of RolePlaying with the same
                 configuration.
         """
+        if self.assistant_agent is None or self.user_agent is None:
+            raise ValueError(
+                "Cannot clone: assistant_agent or user_agent is None"
+            )
 
         new_instance = RolePlaying(
             assistant_role_name=self.assistant_agent.role_name,

docs/key_modules/societies.md

@@ -70,6 +70,8 @@ icon: people-group
         <tr><td>extend_sys_msg_meta_dicts</td><td>List[Dict]</td><td>Extra metadata for system messages</td></tr>
         <tr><td>extend_task_specify_meta_dict</td><td>Dict</td><td>Extra metadata for task specification</td></tr>
         <tr><td>output_language</td><td>str</td><td>Target output language</td></tr>
+        <tr><td>assistant_agent</td><td>ChatAgent</td><td>Custom ChatAgent to use as assistant (optional)</td></tr>
+        <tr><td>user_agent</td><td>ChatAgent</td><td>Custom ChatAgent to use as user (optional)</td></tr>
       </tbody>
     </table>
   </Accordion>