fix: message integration (#3269)

Co-authored-by: Tao Sun <[email protected]>
Co-authored-by: Sun Tao <[email protected]>
Co-authored-by: Wendong-Fan <[email protected]>

Author: 3 people

camel/agents/chat_agent.py

@@ -4519,23 +4519,29 @@ def _clone_tools(
                         # Toolkit doesn't support cloning, use original
                         cloned_toolkits[toolkit_id] = toolkit_instance
 
-                if getattr(
-                    tool.func, "__message_integration_enhanced__", False
-                ):
-                    cloned_tools.append(
-                        FunctionTool(
-                            func=tool.func,
-                            openai_tool_schema=tool.get_openai_tool_schema(),
-                        )
-                    )
-                    continue
-
                 # Get the method from the cloned (or original) toolkit
                 toolkit = cloned_toolkits[toolkit_id]
                 method_name = tool.func.__name__
 
+                # Check if toolkit was actually cloned or just reused
+                toolkit_was_cloned = toolkit is not toolkit_instance
+
                 if hasattr(toolkit, method_name):
                     new_method = getattr(toolkit, method_name)
+
+                    # If toolkit wasn't cloned (stateless), preserve the
+                    # original function to maintain any enhancements/wrappers
+                    if not toolkit_was_cloned:
+                        # Toolkit is stateless, safe to reuse original function
+                        cloned_tools.append(
+                            FunctionTool(
+                                func=tool.func,
+                                openai_tool_schema=tool.get_openai_tool_schema(),
+                            )
+                        )
+                        continue
+
+                    # Toolkit was cloned, use the new method
                     # Wrap cloned method into a new FunctionTool,
                     # preserving schema
                     try:

camel/toolkits/message_integration.py

@@ -148,12 +148,10 @@ def get_message_tool(self) -> FunctionTool:
         """
         return FunctionTool(self.send_message_to_user)
 
-    def register_toolkits(
-        self, toolkit: BaseToolkit, tool_names: Optional[List[str]] = None
-    ) -> BaseToolkit:
-        r"""Add messaging capabilities to toolkit methods.
+    def register_toolkits(self, toolkit: BaseToolkit) -> BaseToolkit:
+        r"""Add messaging capabilities to all toolkit methods.
 
-        This method modifies a toolkit so that specified tools can send
+        This method modifies a toolkit so that all its tools can send
         status messages to users while executing their primary function.
         The tools will accept optional messaging parameters:
         - message_title: Title of the status message
@@ -162,20 +160,18 @@ def register_toolkits(
 
         Args:
             toolkit: The toolkit to add messaging capabilities to
-            tool_names: List of specific tool names to modify.
-                       If None, messaging is added to all tools.
 
         Returns:
-            The toolkit with messaging capabilities added
+            The same toolkit instance with messaging capabilities added to
+                all methods.
         """
         original_tools = toolkit.get_tools()
         enhanced_methods = {}
         for tool in original_tools:
             method_name = tool.func.__name__
-            if tool_names is None or method_name in tool_names:
-                enhanced_func = self._add_messaging_to_tool(tool.func)
-                enhanced_methods[method_name] = enhanced_func
-                setattr(toolkit, method_name, enhanced_func)
+            enhanced_func = self._add_messaging_to_tool(tool.func)
+            enhanced_methods[method_name] = enhanced_func
+            setattr(toolkit, method_name, enhanced_func)
         original_get_tools_method = toolkit.get_tools
 
         def enhanced_get_tools() -> List[FunctionTool]:
@@ -201,7 +197,7 @@ def enhanced_get_tools() -> List[FunctionTool]:
             def enhanced_clone_for_new_session(new_session_id=None):
                 cloned_toolkit = original_clone_method(new_session_id)
                 return message_integration_instance.register_toolkits(
-                    cloned_toolkit, tool_names
+                    cloned_toolkit
                 )
 
             toolkit.clone_for_new_session = enhanced_clone_for_new_session
@@ -300,6 +296,12 @@ def _add_messaging_to_tool(self, func: Callable) -> Callable:
         This internal method modifies the function signature and docstring
         to include optional messaging parameters that trigger status updates.
         """
+        if getattr(func, "__message_integration_enhanced__", False):
+            logger.debug(
+                f"Function {func.__name__} already enhanced, skipping"
+            )
+            return func
+
         # Get the original signature
         original_sig = inspect.signature(func)
 

test/toolkits/test_message_integration.py

@@ -90,16 +90,16 @@ def test_register_all_tools(self):
             self.assertIn('message_attachment', params)
 
     def test_register_specific_tools(self):
-        r"""Test adding messaging to specific tools only."""
-        enhanced_toolkit = self.message_integration.register_toolkits(
-            self.toolkit, tool_names=['search_web']
+        r"""Test adding messaging to specific tools only using
+        register_functions."""
+        # Use register_functions to enhance only one method
+        enhanced_tools = self.message_integration.register_functions(
+            [self.toolkit.search_web]
         )
 
-        tools = enhanced_toolkit.get_tools()
-        search_tool = next(t for t in tools if t.func.__name__ == 'search_web')
-        analyze_tool = next(
-            t for t in tools if t.func.__name__ == 'analyze_data'
-        )
+        # Should get one enhanced tool
+        self.assertEqual(len(enhanced_tools), 1)
+        search_tool = enhanced_tools[0]
 
         # Check search_web has message parameters
         search_schema = search_tool.get_openai_tool_schema()
@@ -108,7 +108,12 @@ def test_register_specific_tools(self):
             search_schema['function']['parameters']['properties'],
         )
 
-        # Check analyze_data doesn't have message parameters
+        # Check the original toolkit's analyze_data doesn't have message
+        # parameters
+        original_tools = self.toolkit.get_tools()
+        analyze_tool = next(
+            t for t in original_tools if t.func.__name__ == 'analyze_data'
+        )
         analyze_schema = analyze_tool.get_openai_tool_schema()
         self.assertNotIn(
             'message_title',