Decorate Agent classes

flaml/autogen/agentchat/__init__.py

@@ -1,5 +1,5 @@
 from .agent import Agent
-from .responsive_agent import ResponsiveAgent
+from .responsive_agent import ResponsiveAgent, register
 from .assistant_agent import AssistantAgent
 from .user_proxy_agent import UserProxyAgent
 from .groupchat import GroupChatManager, GroupChatParticipant
@@ -11,4 +11,5 @@
     "UserProxyAgent",
     "GroupChatManager",
     "GroupChatParticipant",
+    "register",
 ]

flaml/autogen/agentchat/contrib/math_user_proxy_agent.py

@@ -4,7 +4,7 @@
 from typing import Any, Callable, Dict, List, Optional, Union
 from time import sleep
 
-from flaml.autogen.agentchat import Agent, UserProxyAgent
+from flaml.autogen.agentchat import Agent, UserProxyAgent, register
 from flaml.autogen.code_utils import UNKNOWN, extract_code, execute_code, infer_lang
 from flaml.autogen.math_utils import get_answer
 
@@ -123,6 +123,7 @@ def _remove_print(code):
     return "\n".join(lines)
 
 
+@register
 class MathUserProxyAgent(UserProxyAgent):
     """(Experimental) A MathChat agent that can handle math problems."""
 
@@ -165,7 +166,7 @@ def __init__(
             default_auto_reply=default_auto_reply,
             **kwargs,
         )
-        self.register_auto_reply(Agent, self._generate_math_reply)
+        # self.register_auto_reply(Agent, self._generate_math_reply)
         # fixed var
         self._max_invalid_q_per_step = max_invalid_q_per_step
 
@@ -276,7 +277,7 @@ def execute_one_wolfram_query(self, query: str):
             is_success = False
         return output, is_success
 
-    def _generate_math_reply(
+    def _generate_math_reply__auto_reply__(
         self,
         messages: Optional[List[Dict]] = None,
         sender: Optional[Agent] = None,

flaml/autogen/agentchat/responsive_agent.py

@@ -13,6 +13,15 @@ def colored(x, *args, **kwargs):
         return x
 
 
+def register(agent):
+    agent._class_specific_reply = []
+    reply_funcs = [func for func in dir(agent) if func.endswith("__auto_reply__") is True]
+    for reply_func in reply_funcs:
+        agent.register_auto_reply(agent, Agent, reply_func)
+    return agent
+
+
+@register
 class ResponsiveAgent(Agent):
     """(Experimental) A class for generic responsive agents which can be configured as assistant or user proxy.
 
@@ -108,10 +117,16 @@ def __init__(
         self._max_consecutive_auto_reply_dict = defaultdict(self.max_consecutive_auto_reply)
         self._function_map = {} if function_map is None else function_map
         self._default_auto_reply = default_auto_reply
-        self._class_specific_reply = []
-        self.register_auto_reply(Agent, self._generate_oai_reply)
-        self.register_auto_reply(Agent, self._generate_code_execution_reply)
-        self.register_auto_reply(Agent, self._generate_function_call_reply)
+        # self._class_specific_reply = []
+        # print("self._class_specific_reply: ", self._class_specific_reply)
+
+    # def register(class_type):
+    #     def register_reply_func(reply_func: Callable):
+    #         def wrapper(self, *args, **kwargs):
+    #             self.register_auto_reply(class_type, reply_func)
+    #             return reply_func(*args, **kwargs)
+    #         return wrapper
+    #     return register_reply_func
 
     def register_auto_reply(self, class_type, reply_func: Callable):
         """Register a class-specific reply function.
@@ -371,7 +386,7 @@ def clear_history(self, agent: Optional[Agent] = None):
         else:
             self._oai_messages[agent.name].clear()
 
-    def _generate_oai_reply(
+    def _generate_oai_reply__auto_reply__(
         self,
         messages: Optional[List[Dict]] = None,
         sender: Optional[Agent] = None,
@@ -454,7 +469,7 @@ def _check_termination_and_human_reply(
 
         return False, None
 
-    def _generate_function_call_reply(
+    def _generate_function_call_reply__auto_reply__(
         self,
         messages: Optional[List[Dict]] = None,
         sender: Optional[Agent] = None,
@@ -467,7 +482,8 @@ def _generate_function_call_reply(
             return True, func_return
         return False, None
 
-    def _generate_code_execution_reply(
+    # @register(Agent)
+    def _generate_code_execution_reply__auto_reply__(
         self,
         messages: Optional[List[Dict]] = None,
         sender: Optional[Agent] = None,