Continual learning via LearningAgent and TeachingAgent

flaml/autogen/agent/__init__.py

@@ -1,6 +1,9 @@
 from .agent import Agent
 from .assistant_agent import AssistantAgent
 from .user_proxy_agent import UserProxyAgent
+from .teaching_agent import TeachingAgent
+from .learning_agent import LearningAgent
+
 from .math_user_proxy_agent import MathUserProxyAgent
 
-__all__ = ["Agent", "AssistantAgent", "UserProxyAgent", "MathUserProxyAgent"]
+__all__ = ["Agent", "AssistantAgent", "UserProxyAgent", "MathUserProxyAgent", "TeachingAgent", "LearningAgent"]

flaml/autogen/agent/agent.py

@@ -1,5 +1,6 @@
 from collections import defaultdict
 from typing import Dict, Union
+import asyncio
 
 
 class Agent:
@@ -63,11 +64,11 @@ def _append_oai_message(self, message: Union[Dict, str], role, conversation_id):
         oai_message["role"] = "function" if message.get("role") == "function" else role
         self._oai_conversations[conversation_id].append(oai_message)
 
-    def _send(self, message: Union[Dict, str], recipient):
+    async def _send(self, message: Union[Dict, str], recipient):
         """Send a message to another agent."""
         # When the agent composes and sends the message, the role of the message is "assistant". (If 'role' exists and is 'function', it will remain unchanged.)
         self._append_oai_message(message, "assistant", recipient.name)
-        recipient.receive(message, self)
+        await recipient.receive(message, self)
 
     def _receive(self, message: Union[Dict, str], sender):
         """Receive a message from another agent.
@@ -102,6 +103,13 @@ def _receive(self, message: Union[Dict, str], sender):
                     sep="",
                 )
                 print("*" * len(func_print), flush=True)
+
+            # print("message = ",  message, flush=True, sep="")
+            # format the printing of the message
+            print("Message content:", flush=True, sep="")
+            for key, value in message.items():
+                print(f"{key}: {value}", flush=True, sep="\n")
+
         print("\n", "-" * 80, flush=True, sep="")
 
         # When the agent receives a message, the role of the message is "user". (If 'role' exists and is 'function', it will remain unchanged.)

flaml/autogen/agent/learning_agent.py

@@ -0,0 +1,121 @@
+from .assistant_agent import AssistantAgent
+from flaml.autogen.code_utils import DEFAULT_MODEL
+from flaml import oai
+import asyncio
+
+
+class LearningAgent(AssistantAgent):
+    """(Experimental) A learning agent."""
+
+    DEFAULT_SYSTEM_MESSAGE = """You are a helpful AI assistant.
+    In the following cases, suggest python code (in a python coding block) or shell script (in a sh coding block) for the user to execute. You must indicate the script type in the code block.
+    1. When you need to ask the user for some info, use the code to output the info you need, for example, browse or search the web, download/read a file.
+    2. When you need to perform some task with code, use the code to perform the task and output the result. Finish the task smartly. Solve the task step by step if you need to.
+    If you want the user to save the code in a file before executing it, put # filename: <filename> inside the code block as the first line. Don't include multiple code blocks in one response. Do not ask users to copy and paste the result. Instead, use 'print' function for the output when relevant. Check the execution result returned by the user.
+    If the result indicates there is an error, fix the error and output the code again. Suggeset the full code instead of partial code or code changes.
+    Reply "TERMINATE" in the end when everything is done.
+    """
+
+    DEFAULT_CONFIG = {
+        "model": DEFAULT_MODEL,
+    }
+
+    def __init__(self, name, system_message=DEFAULT_SYSTEM_MESSAGE, **config):
+        """
+        Args:
+            name (str): agent name.
+            system_message (str): system message to be sent to the agent.
+            **config (dict): other configurations allowed in
+              [oai.Completion.create](../oai/Completion#create).
+              These configurations will be used when invoking LLM.
+        """
+        super().__init__(name, system_message, **config)
+        self._system_message_learning = """You are a helpful AI assistant."""
+        self._learning_objectives = ""
+        self._can_handle_data_volume = lambda *args: True
+
+    def _generate_task_prompt(self, learning_results, learning_data):
+        """
+        Process the message using NLP.
+        """
+        task_prompt = f"""
+        {self._learning_objectives}.
+        This is the latest data entry: {learning_data}.
+        Renew the current result:
+        {learning_results}
+        You can try to condense the current result and add a new bullet point to the result.
+        """
+        return task_prompt
+
+    @staticmethod
+    def is_total_token_count_within_threshold(learning_results, learning_data):
+        """
+        Check if the total token count of learning data and learning results
+        is within a specified threshold.
+        """
+
+        def _token_counter(input_string):
+            from transformers import GPT2Tokenizer
+
+            # Load a pre-trained tokenizer
+            tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+            # Tokenize the string
+            tokens = tokenizer.tokenize(input_string)
+            return len(tokens)
+
+        oai_max_token_size = 4096
+        return _token_counter(learning_results) + _token_counter(learning_data) < oai_max_token_size * 0.8
+
+    def _validate_learning_constraints(self, learning_constraints):
+        # check if the learning constraints are satisfied
+        # do nothing for now
+        return True
+
+    async def receive(self, message, sender):
+        """Receive a message from another agent."""
+        content = message.get("content", None) if isinstance(message, dict) else message
+        self._receive(message, sender)
+        # NOTE: content and learning settings are mutually exclusive
+        if content is not None:
+            # if content is provided, perform the default receive function
+            super().receive(content, sender)
+        else:
+            # perform learning based on the learning settings
+            learning_func = message.get("learning_func", None)
+            learning_objectives = message.get("learning_objectives", None)
+            learning_constraints = message.get("learning_constraints", None)
+            learning_results = message.get("learning_results", None)
+            data4learning = message.get("data4learning", None)
+            if learning_objectives:
+                self._learning_objectives = learning_objectives
+            # when data is available, perform the learning task when learning_constraints are satisfied
+            if data4learning and self._validate_learning_constraints(learning_constraints):
+                # perform learning
+                if learning_func:
+                    # assumption: learning_func is a function that takes learning_results and learning_data as input and returns new_learning_results and can_handle_data_volume
+                    # when learning_data is None, the learning_func should work as well, outputting the input learning_results as the
+                    # new_learning_results and can_handle_data_volume function
+                    new_learning_results, self._can_handle_data_volume = learning_func(learning_results, data4learning)
+                else:
+                    self._can_handle_data_volume = self.is_total_token_count_within_threshold
+                    if data4learning:
+                        task_prompt = self._generate_task_prompt(learning_results, data4learning)
+                        learning_msg = [
+                            # {"content": self._system_message_learning, "role": "system"},
+                            {"role": "user", "content": task_prompt},
+                        ]
+                        responses = oai.ChatCompletion.create(messages=learning_msg, **self._config)
+                        new_learning_results = oai.ChatCompletion.extract_text(responses)[0]
+                    else:
+                        new_learning_results = learning_results
+                print("*********Current learning results of the learner*********\n", new_learning_results, flush=True)
+                print("*" * 50, flush=True)
+                await self._send(
+                    {"learning_results": new_learning_results, "can_handle_data_volume": self._can_handle_data_volume},
+                    sender,
+                )
+            else:
+                await self._send(
+                    {"learning_results": learning_results, "can_handle_data_volume": self._can_handle_data_volume},
+                    sender,
+                )

flaml/autogen/agent/teaching_agent.py

@@ -0,0 +1,137 @@
+from .user_proxy_agent import UserProxyAgent
+from typing import Optional, Callable
+from transformers import AutoTokenizer
+import asyncio
+
+
+class TeachingAgent(UserProxyAgent):
+    """(Experimental) A teaching agent."""
+
+    def __init__(
+        self,
+        name,
+        system_message="",
+        work_dir=None,
+        human_input_mode="ALWAYS",
+        max_consecutive_auto_reply=None,
+        is_termination_msg=None,
+        use_docker=True,
+        **config,
+    ):
+        """
+        Args:
+            name (str): name of the agent
+            system_message (str): system message to be sent to the agent
+            work_dir (str): working directory for the agent
+            human_input_mode (str): whether to ask for human inputs every time a message is received.
+                Possible values are "ALWAYS", "TERMINATE", "NEVER".
+                (1) When "ALWAYS", the agent prompts for human input every time a message is received.
+                    Under this mode, the conversation stops when the human input is "exit",
+                    or when is_termination_msg is True and there is no human input.
+                (2) When "TERMINATE", the agent only prompts for human input only when a termination message is received or
+                    the number of auto reply reaches the max_consecutive_auto_reply.
+                (3) When "NEVER", the agent will never prompt for human input. Under this mode, the conversation stops
+                    when the number of auto reply reaches the max_consecutive_auto_reply or or when is_termination_msg is True.
+            max_consecutive_auto_reply (int): the maximum number of consecutive auto replies.
+                default to None (no limit provided, class attribute MAX_CONSECUTIVE_AUTO_REPLY will be used as the limit in this case).
+                The limit only plays a role when human_input_mode is not "ALWAYS".
+            is_termination_msg (function): a function that takes a message and returns a boolean value.
+                This function is used to determine if a received message is a termination message.
+            use_docker (bool): whether to use docker to execute the code.
+            **config (dict): other configurations.
+        """
+        super().__init__(
+            name,
+            system_message,
+            work_dir=work_dir,
+            human_input_mode=human_input_mode,
+            max_consecutive_auto_reply=max_consecutive_auto_reply,
+            is_termination_msg=is_termination_msg,
+            use_docker=use_docker,
+            **config,
+        )
+        self._data4learning = []
+        self._learning_constraints = None
+        self._learning_objectives = None
+        self._learning_results = None
+        self._learning_func = None
+        self._can_handle_data_volume = lambda *args: True
+        self._data_available_event = asyncio.Event()
+
+    def setup_learning(
+        self,
+        learning_func: Optional[Callable] = None,
+        learning_objectives: Optional[str] = None,
+        learning_constraints: Optional[dict] = None,
+        learning_results: Optional[str] = "",
+    ):
+        """
+        Args:
+            learning_func (Optional, Callable): the learning function to be executed.
+                The learning function should take the following arguments as inputs:
+                    (1) data4learning: the data for learning.
+                    (2) learning_results: old learning results.
+                The learning function should return the new learning results.
+            learning_objectives (Optional, str): the learning objectives in natural language.
+            learning_constraints (Optional, dict): the learning constraints.
+            learning_results (Optional, str): the learning results in natural language.
+            #TODO: learning_results could be other types of data, e.g., a list of data.
+        Either learning_func or learning_objectives should be provided.
+        """
+        self._learning_constraints = learning_constraints
+        self._learning_objectives = learning_objectives  # already reflected in the learning_func
+        self._learning_results = learning_results
+        self._learning_func = learning_func
+        assert (
+            self._learning_func is not None or self._learning_objectives is not None
+        ), "learning_func or learning_objectives should be provided"
+
+        self._learning_settings = {
+            "learning_func": self._learning_func,
+            "learning_objectives": self._learning_objectives,
+            "learning_constraints": self._learning_constraints,
+            "learning_results": self._learning_results,
+            "data4learning": [],
+        }
+
+    def generate_init_prompt(self):
+        """
+        When generating the init prompt, we need to distinguish the two cases where learning_func or learning_objectives is provided.
+        """
+        self._init_prompt = self._learning_settings.copy()
+
+        return self._init_prompt
+
+    async def add_data(self, data4learning):
+        """Add data for learning."""
+        self._data4learning += data4learning
+        print(f"{len(data4learning)} data entries added for learning!")
+        self._data_available_event.set()
+
+    async def auto_reply(self, message, sender, default_reply=""):
+        """
+        Need to distinguish if the sender is requesting for learning data or not
+        """
+        learning_results = message.get("learning_results", "")
+        can_handle_data_volume = message.get("can_handle_data_volume") or self._can_handle_data_volume
+        current_data4learning = []
+        # Wait here if no data is available
+        while not self._data4learning:
+            print("waiting for data...")
+            await self._data_available_event.wait()
+        # Reset the event as we are going to consume data
+        self._data_available_event.clear()
+        while self._data4learning:
+            combined_data_str = "\n".join(current_data4learning + [self._data4learning[0]])
+            if can_handle_data_volume(learning_results, combined_data_str):
+                current_data4learning.append(self._data4learning.pop(0))
+            else:
+                break
+        if current_data4learning:
+            response = {
+                "learning_results": learning_results,
+                "data4learning": current_data4learning,
+            }
+            await self._send(response, sender)
+        else:
+            print("no data for learning and thus terminate the conversation")

flaml/autogen/agent/user_proxy_agent.py

@@ -3,6 +3,7 @@
 from collections import defaultdict
 import json
 from typing import Dict, Union
+import asyncio
 
 
 class UserProxyAgent(Agent):
@@ -173,24 +174,24 @@ def _execute_function(self, func_call):
             "content": str(content),
         }
 
-    def auto_reply(self, message: dict, sender, default_reply=""):
+    async def auto_reply(self, message: dict, sender, default_reply=""):
         """Generate an auto reply."""
         if "function_call" in message:
             is_exec_success, func_return = self._execute_function(message["function_call"])
-            self._send(func_return, sender)
+            await self._send(func_return, sender)
             return
 
         code_blocks = extract_code(message["content"])
         if len(code_blocks) == 1 and code_blocks[0][0] == UNKNOWN:
             # no code block is found, lang should be `UNKNOWN`
-            self._send(default_reply, sender)
+            await self._send(default_reply, sender)
         else:
             # try to execute the code
             exitcode, logs = self._execute_code(code_blocks)
             exitcode2str = "execution succeeded" if exitcode == 0 else "execution failed"
-            self._send(f"exitcode: {exitcode} ({exitcode2str})\nCode output: {logs}", sender)
+            await self._send(f"exitcode: {exitcode} ({exitcode2str})\nCode output: {logs}", sender)
 
-    def receive(self, message: Union[Dict, str], sender):
+    async def receive(self, message: Union[Dict, str], sender):
         """Receive a message from the sender agent.
         Once a message is received, this function sends a reply to the sender or simply stop.
         The reply can be generated automatically or entered manually by a human.
@@ -221,9 +222,9 @@ def receive(self, message: Union[Dict, str], sender):
         if reply:
             # reset the consecutive_auto_reply_counter
             self._consecutive_auto_reply_counter[sender.name] = 0
-            self._send(reply, sender)
+            await self._send(reply, sender)
             return
 
         self._consecutive_auto_reply_counter[sender.name] += 1
         print("\n>>>>>>>> NO HUMAN INPUT RECEIVED. USING AUTO REPLY FOR THE USER...", flush=True)
-        self.auto_reply(message, sender, default_reply=reply)
+        await self.auto_reply(message, sender, default_reply=reply)