Repository: microsoft/FLAML (language: python)

Issue:
`infer_lang` in autogen.code_utils
https://github.com/microsoft/FLAML/blob/main/flaml/autogen/code_utils.py#L24

sometimes the sh code would start with "python3 ", so maybe include "python3 " as well?


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/openai.yml b/.github/workflows/openai.yml
index 11b9dcd988..f1949bc9f0 100644
--- a/.github/workflows/openai.yml
+++ b/.github/workflows/openai.yml
@@ -34,6 +34,10 @@ jobs:
           pip install -e .[autogen,blendsearch]
           python -c "import flaml"
           pip install coverage pytest datasets
+      - name: Install packages for test when needed
+        if: matrix.python-version == '3.9'
+        run: |
+          pip install -e .[optiguide] docker
       - name: Install packages for MathChat when needed
         if: matrix.python-version != '3.11'
         run: |
diff --git a/README.md b/README.md
index 73bc50e7b2..a2a1d58c8a 100644
--- a/README.md
+++ b/README.md
@@ -66,10 +66,10 @@ Use the following guides to get started with FLAML in .NET:
 * (New) The [autogen](https://microsoft.github.io/FLAML/docs/Use-Cases/Auto-Generation) package can help you maximize the utility out of the expensive LLMs such as ChatGPT and GPT-4, including:
     - A drop-in replacement of `openai.Completion` or `openai.ChatCompletion` with powerful functionalites like tuning, caching, templating, filtering. For example, you can optimize generations by LLM with your own tuning data, success metrics and budgets.
     ```python
-    from flaml import oai
+    from flaml import autogen
 
     # perform tuning
-    config, analysis = oai.Completion.tune(
+    config, analysis = autogen.Completion.tune(
         data=tune_data,
         metric="success",
         mode="max",
@@ -80,13 +80,13 @@ Use the following guides to get started with FLAML in .NET:
     )
 
     # perform inference for a test instance
-    response = oai.Completion.create(context=test_instance, **config)
+    response = autogen.Completion.create(context=test_instance, **config)
     ```
     - LLM-driven intelligent agents which can collaborately perform tasks autonomously or with human feedback, including tasks that require using tools via code.
     ```python
-    assistant = AssistantAgent("assistant")
-    user_proxy = UserProxyAgent("user_proxy")
-    user_proxy.initiate_chat("Show me the YTD gain of 10 largest technology companies as of today.")
+    assistant = autogen.AssistantAgent("assistant")
+    user_proxy = autogen.UserProxyAgent("user_proxy")
+    user_proxy.initiate_chat(assistant, message="Show me the YTD gain of 10 largest technology companies as of today.")
     ```
 * With three lines of code, you can start using this economical and fast
 AutoML engine as a [scikit-learn style estimator](https://microsoft.github.io/FLAML/docs/Use-Cases/Task-Oriented-AutoML).
diff --git a/flaml/__init__.py b/flaml/__init__.py
index 9fca486499..2349993733 100644
--- a/flaml/__init__.py
+++ b/flaml/__init__.py
@@ -2,7 +2,6 @@
 from flaml.automl import AutoML, logger_formatter
 from flaml.tune.searcher import CFO, BlendSearch, FLOW2, BlendSearchTuner, RandomSearch
 from flaml.onlineml.autovw import AutoVW
-from flaml.autogen import oai
 from flaml.version import __version__
 
 
diff --git a/flaml/autogen/__init__.py b/flaml/autogen/__init__.py
index d13c36540e..949f7d3ce8 100644
--- a/flaml/autogen/__init__.py
+++ b/flaml/autogen/__init__.py
@@ -1,2 +1,3 @@
-DEFAULT_MODEL = "gpt-4"
-FAST_MODEL = "gpt-3.5-turbo"
+from .oai import *
+from .agentchat import *
+from .code_utils import DEFAULT_MODEL, FAST_MODEL
diff --git a/flaml/autogen/agentchat/agent.py b/flaml/autogen/agentchat/agent.py
index c62ed29ab4..27c552d1f5 100644
--- a/flaml/autogen/agentchat/agent.py
+++ b/flaml/autogen/agentchat/agent.py
@@ -38,7 +38,7 @@ def generate_reply(
         messages: Optional[List[Dict]] = None,
         default_reply: Optional[Union[str, Dict]] = "",
         sender: Optional["Agent"] = None,
-    ) -> Union[str, Dict]:
+    ) -> Union[str, Dict, None]:
         """(Abstract method) Generate a reply based on the received messages.
 
         Args:
@@ -46,5 +46,5 @@ def generate_reply(
             default_reply (str or dict): the default reply if no other reply is generated.
             sender: sender of an Agent instance.
         Returns:
-            str or dict: the generated reply.
+            str or dict or None: the generated reply. If None, no reply is generated.
         """
diff --git a/flaml/autogen/agentchat/assistant_agent.py b/flaml/autogen/agentchat/assistant_agent.py
index 00f027edb7..f8e48a0f9f 100644
--- a/flaml/autogen/agentchat/assistant_agent.py
+++ b/flaml/autogen/agentchat/assistant_agent.py
@@ -27,7 +27,7 @@ def __init__(
         self,
         name: str,
         system_message: Optional[str] = DEFAULT_SYSTEM_MESSAGE,
-        oai_config: Optional[Union[Dict, bool]] = None,
+        llm_config: Optional[Union[Dict, bool]] = None,
         is_termination_msg: Optional[Callable[[Dict], bool]] = None,
         max_consecutive_auto_reply: Optional[int] = None,
         human_input_mode: Optional[str] = "NEVER",
@@ -37,10 +37,10 @@ def __init__(
         """
         Args:
             name (str): agent name.
-            system_message (str): system message for the oai inference.
+            system_message (str): system message for the ChatCompletion inference.
                 Please override this attribute if you want to reprogram the agent.
-            oai_config (dict): oai inference configuration.
-                Please refer to [oai.Completion.create](/docs/reference/autogen/oai/completion#create)
+            llm_config (dict): llm inference configuration.
+                Please refer to [autogen.Completion.create](/docs/reference/autogen/oai/completion#create)
                 for available options.
             is_termination_msg (function): a function that takes a message in the form of a dictionary
                 and returns a boolean value indicating if this received message is a termination message.
@@ -58,6 +58,6 @@ def __init__(
             max_consecutive_auto_reply,
             human_input_mode,
             code_execution_config=code_execution_config,
-            oai_config=oai_config,
+            llm_config=llm_config,
             **kwargs,
         )
diff --git a/flaml/autogen/agentchat/contrib/math_user_proxy_agent.py b/flaml/autogen/agentchat/contrib/math_user_proxy_agent.py
index 2533f9be41..3e90fbcf04 100644
--- a/flaml/autogen/agentchat/contrib/math_user_proxy_agent.py
+++ b/flaml/autogen/agentchat/contrib/math_user_proxy_agent.py
@@ -220,14 +220,6 @@ def execute_one_python_code(self, pycode):
         return_code, output, _ = execute_code(pycode, **self._code_execution_config, timeout=5)
         is_success = return_code == 0
 
-        # Decode the output
-        if isinstance(output, bytes):
-            try:
-                output = output.decode("utf-8")
-            except UnicodeDecodeError:
-                is_success = False
-                output = "The return cannot be decoded."
-
         if not is_success:
             # Remove the file information from the error string
             pattern = r'File "/[^"]+\.py", line \d+, in .+\n'
@@ -285,11 +277,11 @@ def generate_reply(
         self,
         messages: Optional[List[Dict]] = None,
         default_reply: Optional[Union[str, Dict]] = DEFAULT_REPLY,
-        sender: Optional["Agent"] = None,
-    ) -> Union[str, Dict]:
+        sender: Optional[Agent] = None,
+    ) -> Union[str, Dict, None]:
         """Generate an auto reply."""
         if messages is None:
-            messages = self._oai_conversations[sender.name]
+            messages = self._oai_messages[sender.name]
         message = messages[-1]
         message = message.get("content", "")
         code_blocks = extract_code(message)
diff --git a/flaml/autogen/agentchat/contrib/opti_guide.py b/flaml/autogen/agentchat/contrib/opti_guide.py
index 6e0208ce1f..036a86a512 100644
--- a/flaml/autogen/agentchat/contrib/opti_guide.py
+++ b/flaml/autogen/agentchat/contrib/opti_guide.py
@@ -4,7 +4,7 @@
 The design here is that a user asked a question, then OptiGuide will answer
 it.
 
-The OptiGuide agent will interact with the LLMProxy agent to access LLMs.
+The OptiGuide agent will interact with LLM-based agents.
 
 Notes:
 1. We assume there is a Gurobi model `m` in the global scope.
@@ -14,42 +14,34 @@
 4. We simplify the evaluation only to "DATA CODE" and "CONSTRAINT CODE", where
 we would insert the newly added code.
 """
-import pdb
 import re
-from typing import Dict, Union
+from typing import Dict, List, Optional, Union
+from flaml.autogen.agentchat.agent import Agent
 
-import gurobipy as gp
 from eventlet.timeout import Timeout
 from gurobipy import GRB
 from termcolor import colored
 
-from flaml import oai
-from flaml.autogen.agentchat import AssistantAgent, ResponsiveAgent, UserProxyAgent
-from flaml.autogen.code_utils import DEFAULT_MODEL, extract_code
+from flaml.autogen.agentchat import AssistantAgent
+from flaml.autogen.code_utils import extract_code
 
 # %% System Messages
-ASSIST_SYSTEM_MSG = """You are OptiGuide,
-an agent to write Python code and to answer users questions for supply chain-related coding project.
-
-
-You are a chatbot to explain solutions from a Gurobi/Python solver.
-
+PENCIL_SYSTEM_MSG = """You are a chatbot to:
+(1) write Python code to answer users questions for supply chain-related coding project;
+(2) explain solutions from a Gurobi/Python solver.
 
 --- SOURCE CODE ---
 {source_code}
 
 --- DOC STR ---
 {doc_str}
-
+---
 
 Here are some example questions and their answers and codes:
 --- EXAMPLES ---
 {example_qa}
-
-
 ---
 
-
 The execution result of the original source code is below.
 --- Original Result ---
 {execution_result}
@@ -59,7 +51,7 @@
 You just need to write code snippet.
 """
 
-SAFEGUARD_SYSTEM_MSG = """
+SHIELD_SYSTEM_MSG = """
 Given the original source code:
 {source_code}
 
@@ -74,130 +66,106 @@
 
 
 # %%
-class OptiGuideAgent(ResponsiveAgent):
-    """(Experimental) OptiGuide is an agent to write Python code and to answer
-      users questions for supply chain-related coding project.
+class OptiGuideAgent(AssistantAgent):
+    """(Experimental) OptiGuide is an agent to answer
+    users questions for supply chain-related coding project.
 
-    Here, the OptiGuide agent manages three agents (coder, safeguard, and interpreter)
-    and two assistant agents (pencil and shield).
+    Here, the OptiGuide agent manages two assistant agents (pencil and shield).
     """
 
-    def __init__(self, name, source_code, doc_str="", example_qa="", **config):
+    def __init__(self, name, source_code, doc_str="", example_qa="", debug_times=3, **config):
         """
         Args:
             name (str): agent name.
             source_code (str): The original source code to run.
             doc_str (str): docstring for helper functions if existed.
             example_qa (str): training examples for in-context learning.
+            debug_times (int): number of debug tries we allow for LLM.
             **config (dict): other configurations allowed in
-              [oai.Completion.create](../oai/Completion#create).
-              These configurations will be used when invoking LLM.
+              [ResponsiveAgent](../responsive_agent/ResponsiveAgent#__init__).
         """
         assert source_code.find(DATA_CODE_STR) >= 0, "DATA_CODE_STR not found."
         assert source_code.find(CONSTRAINT_CODE_STR) >= 0, "CONSTRAINT_CODE_STR not found."
 
-        super().__init__(
-            name, max_consecutive_auto_reply=0, human_input_mode="NEVER", code_execution_config=False, **config
-        )
+        super().__init__(name, **config)
         self._source_code = source_code
         self._doc_str = doc_str
         self._example_qa = example_qa
         self._origin_execution_result = _run_with_exec(source_code)
-
-        self._debug_opportunity = 3  # number of debug tries we allow for LLM
-
-    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.
-
-        OptiGuide will invoke its components: coder, safeguard, and interpreter
-        to generate the reply.
-
-        The steps follow the one in Figure 2 of the OptiGuide paper.
-        """
-        # Step 1: receive the message
-        message = self._message_to_dict(message)
-        super().receive(message, sender)
-        # default reply is empty (i.e., no reply, in this case we will try to generate auto reply)
-        reply = ""
-        message["content"]
-
-        user_chat_history = f"\nHere are the history of discussions:\n{self._oai_conversations[sender.name]}"
-
-        # Spawn coder, interpreter, and safeguard
-        # Note: we use the same agent for coding and interpreting, so that they
-        # can share the same chat history.
-        coder = interpreter = UserProxyAgent(
-            "coder interpreter", human_input_mode="NEVER", max_consecutive_auto_reply=0
-        )
-        safeguard = UserProxyAgent("safeguard", human_input_mode="NEVER", max_consecutive_auto_reply=0)
-
-        # Spawn the assistants for coder, interpreter, and safeguard
-        writing_sys_msg = ASSIST_SYSTEM_MSG.format(
-            source_code=self._source_code,
-            doc_str=self._doc_str,
-            example_qa=self._example_qa,
-            execution_result=self._origin_execution_result,
-        )
-        shield_sys_msg = SAFEGUARD_SYSTEM_MSG.format(source_code=self._source_code)
-        pencil = AssistantAgent("pencil", system_message=writing_sys_msg + user_chat_history)
-        shield = AssistantAgent("shield", system_message=shield_sys_msg + user_chat_history)
-
-        debug_opportunity = self._debug_opportunity
-
-        new_code = None
-        execution_rst = None
-        success = False
-        while execution_rst is None and debug_opportunity > 0:
-            # Step 2: Write code
-            if new_code is None:
-                # The first time we are trying to solve the question in this session
-                msg = CODE_PROMPT
-            else:
-                # debug code
-                msg = DEBUG_PROMPT.format(execution_rst)
-
-            coder.send(message=msg, recipient=pencil)
-            new_code = coder.oai_conversations[pencil.name][-1]["content"]
-            new_code = extract_code(new_code)[0][1]  # First code block, the code (excluding language)
-            print(colored(new_code, "green"))
-
-            # Step 3: safeguard
-            # FYI: in production, there are some other steps to check the code; for instance,
-            # with a static code analyzer or some other external packages.
-            safeguard.send(message=SAFEGUARD_PROMPT.format(code=new_code), recipient=shield)
-            safe_msg = safeguard.oai_conversations[shield.name][-1]["content"]
-            print("Safety:", colored(safe_msg, "blue"))
-
-            if safe_msg.find("DANGER") < 0:
-                # Step 4 and 5: Run the code and obtain the results
-                src_code = _insert_code(self._source_code, new_code)
-                execution_rst = _run_with_exec(src_code)
-
-                print(colored(str(execution_rst), "yellow"))
-                if type(execution_rst) in [str, int, float]:
-                    success = True
-                    break  # we successfully run the code and get the result
+        self._pencil = AssistantAgent("pencil", llm_config=self.llm_config)
+        self._shield = AssistantAgent("shield", llm_config=self.llm_config)
+        self._debug_times_left = self.debug_times = 3
+        self._success = False
+
+    def generate_reply(
+        self,
+        messages: Optional[List[Dict]] = None,
+        default_reply: Optional[Union[str, Dict]] = "",
+        sender: Optional[Agent] = None,
+    ) -> Union[str, Dict, None]:
+        """Reply based on the conversation history."""
+        if sender not in [self._pencil, self._shield]:
+            # Step 1: receive the message from the user
+            user_chat_history = f"\nHere are the history of discussions:\n{self._oai_messages[sender.name]}"
+            pencil_sys_msg = (
+                PENCIL_SYSTEM_MSG.format(
+                    source_code=self._source_code,
+                    doc_str=self._doc_str,
+                    example_qa=self._example_qa,
+                    execution_result=self._origin_execution_result,
+                )
+                + user_chat_history
+            )
+            shield_sys_msg = SHIELD_SYSTEM_MSG.format(source_code=self._source_code) + user_chat_history
+            self._pencil.update_system_message(pencil_sys_msg)
+            self._shield.update_system_message(shield_sys_msg)
+            self._pencil.reset()
+            self._shield.reset()
+            self._debug_times_left = self.debug_times
+            self._success = False
+            # Step 2-6: code, safeguard, and interpret
+            self.initiate_chat(self._pencil, message=CODE_PROMPT)
+            if self._success:
+                # step 7: receive interpret result
+                reply = self.last_message(self._pencil)["content"]
             else:
-                # DANGER: If not safe, try to debug. Redo coding
-                execution_rst = """
-                Sorry, this new code is not safe to run. I would not allow you to execute it.
-                Please try to find a new way (coding) to answer the question."""
-
-            # Try to debug and write code again
-            debug_opportunity -= 1
-
-        if success:
-            # Step 6 - 7: interpret results
-            interpret_msg = INTERPRETER_PROMPT.format(execution_rst=execution_rst)
-            interpreter.send(message=interpret_msg, recipient=pencil)
-            reply = interpreter.oai_conversations[pencil.name][-1]["content"]
+                reply = "Sorry. I cannot answer your question."
+            # Finally, step 8: send reply to user
+            return reply
+        if sender == self._pencil:
+            # reply to pencil
+            return self._generate_reply_to_pencil(sender)
+        # no reply to shield
+
+    def _generate_reply_to_pencil(self, sender):
+        if self._success:
+            # no reply to pencil
+            return
+        # Step 3: safeguard
+        _, code = extract_code(self.last_message(sender)["content"])[0]
+        # print(colored(code, "green"))
+        self.initiate_chat(message=SAFEGUARD_PROMPT.format(code=code), recipient=self._shield)
+        safe_msg = self.last_message(self._shield)["content"]
+        # print("Safety:", colored(safe_msg, "blue"))
+        if safe_msg.find("DANGER") < 0:
+            # Step 4 and 5: Run the code and obtain the results
+            src_code = _insert_code(self._source_code, code)
+            execution_rst = _run_with_exec(src_code)
+            # print(colored(str(execution_rst), "yellow"))
+            if type(execution_rst) in [str, int, float]:
+                # we successfully run the code and get the result
+                self._success = True
+                # Step 6: request to interpret results
+                return INTERPRETER_PROMPT.format(execution_rst=execution_rst)
         else:
-            reply = "Sorry. I cannot answer your question."
-
-        # Finally, step 8: send reply
-        self.send(reply, sender)
+            # DANGER: If not safe, try to debug. Redo coding
+            execution_rst = """
+Sorry, this new code is not safe to run. I would not allow you to execute it.
+Please try to find a new way (coding) to answer the question."""
+        if self._debug_times_left > 0:
+            # Try to debug and write code again (back to step 2)
+            self._debug_times_left -= 1
+            return DEBUG_PROMPT.format(error_message=execution_rst)
 
 
 # %% Helper functions to edit and run code.
diff --git a/flaml/autogen/agentchat/responsive_agent.py b/flaml/autogen/agentchat/responsive_agent.py
index c240052dcb..43e58b27f1 100644
--- a/flaml/autogen/agentchat/responsive_agent.py
+++ b/flaml/autogen/agentchat/responsive_agent.py
@@ -1,10 +1,17 @@
 from collections import defaultdict
 import json
 from typing import Callable, Dict, List, Optional, Union
-from flaml import oai
+from flaml.autogen import oai
 from .agent import Agent
 from flaml.autogen.code_utils import DEFAULT_MODEL, UNKNOWN, execute_code, extract_code, infer_lang
 
+try:
+    from termcolor import colored
+except ImportError:
+
+    def colored(x, *args, **kwargs):
+        return x
+
 
 class ResponsiveAgent(Agent):
     """(Experimental) A class for generic responsive agents which can be configured as assistant or user proxy.
@@ -35,12 +42,12 @@ def __init__(
         human_input_mode: Optional[str] = "TERMINATE",
         function_map: Optional[Dict[str, Callable]] = None,
         code_execution_config: Optional[Union[Dict, bool]] = None,
-        oai_config: Optional[Union[Dict, bool]] = None,
+        llm_config: Optional[Union[Dict, bool]] = None,
     ):
         """
         Args:
             name (str): name of the agent.
-            system_message (str): system message for the oai inference.
+            system_message (str): system message for the ChatCompletion inference.
             is_termination_msg (function): a function that takes a message in the form of a dictionary
                 and returns a boolean value indicating if this received message is a termination message.
                 The dict can contain the following keys: "content", "role", "name", "function_call".
@@ -71,25 +78,24 @@ def __init__(
                     If the code is executed in the current environment,
                     the code must be trusted.
                 - timeout (Optional, int): The maximum execution time in seconds.
-            oai_config (dict or False): oai inference configuration.
-                Please refer to [oai.Completion.create](/docs/reference/autogen/oai/completion#create)
+            llm_config (dict or False): llm inference configuration.
+                Please refer to [autogen.Completion.create](/docs/reference/autogen/oai/completion#create)
                 for available options.
-                To disable oai-based auto reply, set to False.
+                To disable llm-based auto reply, set to False.
         """
         super().__init__(name)
         # a dictionary of conversations, default value is list
-        self._oai_conversations = defaultdict(list)
-        self._system_message = system_message
-        self._oai_system_message = [{"content": self._system_message, "role": "system"}]
+        self._oai_messages = defaultdict(list)
+        self._oai_system_message = [{"content": system_message, "role": "system"}]
         self._is_termination_msg = (
             is_termination_msg if is_termination_msg is not None else (lambda x: x.get("content") == "TERMINATE")
         )
-        if oai_config is False:
-            self.oai_config = False
+        if llm_config is False:
+            self.llm_config = False
         else:
-            self.oai_config = self.DEFAULT_CONFIG.copy()
-            if isinstance(oai_config, dict):
-                self.oai_config.update(oai_config)
+            self.llm_config = self.DEFAULT_CONFIG.copy()
+            if isinstance(llm_config, dict):
+                self.llm_config.update(llm_config)
 
         self._code_execution_config = {} if code_execution_config is None else code_execution_config
         self.human_input_mode = human_input_mode
@@ -99,10 +105,18 @@ def __init__(
         self._consecutive_auto_reply_counter = defaultdict(int)
         self._function_map = {} if function_map is None else function_map
 
+    def update_system_message(self, system_message: str):
+        """Update the system message.
+
+        Args:
+            system_message (str): system message for the ChatCompletion inference.
+        """
+        self._oai_system_message[0]["content"] = system_message
+
     @property
-    def oai_conversations(self) -> Dict[str, List[Dict]]:
-        """A dictionary of conversations from name to list of oai messages."""
-        return self._oai_conversations
+    def chat_messages(self) -> Dict[str, List[Dict]]:
+        """A dictionary of conversations from name to list of ChatCompletion messages."""
+        return self._oai_messages
 
     def last_message(self, agent: Optional[Agent] = None) -> Dict:
         """The last message exchanged with the agent.
@@ -116,14 +130,14 @@ def last_message(self, agent: Optional[Agent] = None) -> Dict:
             The last message exchanged with the agent.
         """
         if agent is None:
-            n_conversations = len(self._oai_conversations)
+            n_conversations = len(self._oai_messages)
             if n_conversations == 0:
                 return None
             if n_conversations == 1:
-                for conversation in self._oai_conversations.values():
+                for conversation in self._oai_messages.values():
                     return conversation[-1]
             raise ValueError("More than one conversation is found. Please specify the sender to get the last message.")
-        return self._oai_conversations[agent.name][-1]
+        return self._oai_messages[agent.name][-1]
 
     @property
     def use_docker(self) -> Union[bool, str, None]:
@@ -143,19 +157,19 @@ def _message_to_dict(message: Union[Dict, str]):
             return message
 
     def _append_oai_message(self, message: Union[Dict, str], role, conversation_id) -> bool:
-        """Append a message to the oai conversation.
+        """Append a message to the ChatCompletion conversation.
 
         If the message received is a string, it will be put in the "content" field of the new dictionary.
         If the message received is a dictionary but does not have any of the two fields "content" or "function_call",
-            this message is not a valid oai message and will be ignored.
+            this message is not a valid ChatCompletion message.
 
         Args:
-            message (dict or str): message to be appended to the oai conversation.
+            message (dict or str): message to be appended to the ChatCompletion conversation.
             role (str): role of the message, can be "assistant" or "function".
             conversation_id (str): id of the conversation, should be the name of the recipient or sender.
 
         Returns:
-            bool: whether the message is appended to the oai conversation.
+            bool: whether the message is appended to the ChatCompletion conversation.
         """
         message = self._message_to_dict(message)
         # create oai message to be appended to the oai conversation that can be passed to oai directly.
@@ -164,10 +178,10 @@ def _append_oai_message(self, message: Union[Dict, str], role, conversation_id)
             return False
 
         oai_message["role"] = "function" if message.get("role") == "function" else role
-        self._oai_conversations[conversation_id].append(oai_message)
+        self._oai_messages[conversation_id].append(oai_message)
         return True
 
-    def send(self, message: Union[Dict, str], recipient: "Agent"):
+    def send(self, message: Union[Dict, str], recipient: Agent):
         """Send a message to another agent.
 
         Args:
@@ -179,16 +193,16 @@ def send(self, message: Union[Dict, str], recipient: "Agent"):
                 - role (str): the role of the message, any role that is not "function"
                     will be modified to "assistant".
                 - context (dict): the context of the message, which will be passed to
-                    [oai.Completion.create](../oai/Completion#create).
+                    [autogen.Completion.create](../oai/Completion#create).
                     For example, one agent can send a message A as:
-            ```python
-            {
-                "content": "{use_tool_msg}",
-                "context": {
-                    "use_tool_msg": "Use tool X if they are relevant."
-                }
+        ```python
+        {
+            "content": "{use_tool_msg}",
+            "context": {
+                "use_tool_msg": "Use tool X if they are relevant."
             }
-            ```
+        }
+        ```
                     Next time, one agent can send a message B with a different "use_tool_msg".
                     Then the content of message A will be refreshed to the new "use_tool_msg".
                     So effectively, this provides a way for an agent to send a "link" and modify
@@ -196,7 +210,7 @@ def send(self, message: Union[Dict, str], recipient: "Agent"):
             recipient (Agent): the recipient of the message.
 
         Raises:
-            ValueError: if the message is not a valid oai message.
+            ValueError: if the message can't be converted into a valid ChatCompletion message.
         """
         # When the agent composes and sends the message, the role of the message is "assistant"
         # unless it's "function".
@@ -204,51 +218,58 @@ def send(self, message: Union[Dict, str], recipient: "Agent"):
         if valid:
             recipient.receive(message, self)
         else:
-            raise ValueError("Message is not a valid oai message. Either content or function_call must be provided.")
+            raise ValueError(
+                "Message can't be converted into a valid ChatCompletion message. Either content or function_call must be provided."
+            )
 
-    def _print_received_message(self, message: Union[Dict, str], sender: "Agent"):
+    def _print_received_message(self, message: Union[Dict, str], sender: Agent):
         # print the message received
-        print(sender.name, "(to", f"{self.name}):\n", flush=True)
+        print(colored(sender.name, "yellow"), "(to", f"{self.name}):\n", flush=True)
         if message.get("role") == "function":
             func_print = f"***** Response from calling function \"{message['name']}\" *****"
-            print(func_print, flush=True)
+            print(colored(func_print, "green"), flush=True)
             print(message["content"], flush=True)
-            print("*" * len(func_print), flush=True)
+            print(colored("*" * len(func_print), "green"), flush=True)
         else:
             if message.get("content") is not None:
                 print(message["content"], flush=True)
             if "function_call" in message:
                 func_print = f"***** Suggested function Call: {message['function_call'].get('name', '(No function name found)')} *****"
-                print(func_print, flush=True)
+                print(colored(func_print, "green"), flush=True)
                 print(
                     "Arguments: \n",
                     message["function_call"].get("arguments", "(No arguments found)"),
                     flush=True,
                     sep="",
                 )
-                print("*" * len(func_print), flush=True)
+                print(colored("*" * len(func_print), "green"), flush=True)
         print("\n", "-" * 80, flush=True, sep="")
 
-    def receive(self, message: Union[Dict, str], sender: "Agent"):
+    def receive(self, message: Union[Dict, str], sender: Agent):
         """Receive a message from another agent.
 
         Once a message is received, this function sends a reply to the sender or stop.
         The reply can be generated automatically or entered manually by a human.
 
         Args:
-            message (dict or str): message from the sender. If the type is dict, it may contain the following reserved fields (All fields are optional).
+            message (dict or str): message from the sender. If the type is dict, it may contain the following reserved fields (either content or function_call need to be provided).
                 1. "content": content of the message, can be None.
                 2. "function_call": a dictionary containing the function name and arguments.
                 3. "role": role of the message, can be "assistant", "user", "function".
                     This field is only needed to distinguish between "function" or "assistant"/"user".
                 4. "name": In most cases, this field is not needed. When the role is "function", this field is needed to indicate the function name.
             sender: sender of an Agent instance.
+
+        Raises:
+            ValueError: if the message can't be converted into a valid ChatCompletion message.
         """
         message = self._message_to_dict(message)
         # When the agent receives a message, the role of the message is "user". (If 'role' exists and is 'function', it will remain unchanged.)
         valid = self._append_oai_message(message, "user", sender.name)
         if not valid:
-            return
+            raise ValueError(
+                "Received message can't be converted into a valid ChatCompletion message. Either content or function_call must be provided."
+            )
         self._print_received_message(message, sender)
 
         # default reply is empty (i.e., no reply, in this case we will try to generate auto reply)
@@ -290,7 +311,7 @@ def receive(self, message: Union[Dict, str], sender: "Agent"):
 
         # print the no_human_input_msg
         if no_human_input_msg:
-            print(f"\n>>>>>>>> {no_human_input_msg}", flush=True)
+            print(colored(f"\n>>>>>>>> {no_human_input_msg}", "red"), flush=True)
 
         # stop the conversation
         if reply == "exit":
@@ -308,18 +329,20 @@ def receive(self, message: Union[Dict, str], sender: "Agent"):
         # send the auto reply
         self._consecutive_auto_reply_counter[sender.name] += 1
         if self.human_input_mode != "NEVER":
-            print("\n>>>>>>>> USING AUTO REPLY...", flush=True)
-        self.send(self.generate_reply(sender=sender), sender)
+            print(colored("\n>>>>>>>> USING AUTO REPLY...", "red"), flush=True)
+        reply = self.generate_reply(sender=sender)
+        if reply is not None:
+            self.send(reply, sender)
 
     def reset(self):
         """Reset the agent."""
-        self._oai_conversations.clear()
+        self._oai_messages.clear()
         self._consecutive_auto_reply_counter.clear()
 
     def _oai_reply(self, messages: List[Dict]) -> Union[str, Dict]:
         # TODO: #1143 handle token limit exceeded error
         response = oai.ChatCompletion.create(
-            context=messages[-1].get("context"), messages=self._oai_system_message + messages, **self.oai_config
+            context=messages[-1].get("context"), messages=self._oai_system_message + messages, **self.llm_config
         )
         return oai.ChatCompletion.extract_text_or_function_call(response)[0]
 
@@ -327,8 +350,8 @@ def generate_reply(
         self,
         messages: Optional[List[Dict]] = None,
         default_reply: Optional[Union[str, Dict]] = "",
-        sender: Optional["Agent"] = None,
-    ) -> Union[str, Dict]:
+        sender: Optional[Agent] = None,
+    ) -> Union[str, Dict, None]:
         """Reply based on the conversation history.
 
         First, execute function or code and return the result.
@@ -342,21 +365,21 @@ def generate_reply(
             sender: sender of an Agent instance.
 
         Returns:
-            str or dict: reply.
+            str or dict or None: reply. None if no reply is generated.
         """
         assert messages is not None or sender is not None, "Either messages or sender must be provided."
         if messages is None:
-            messages = self._oai_conversations[sender.name]
+            messages = self._oai_messages[sender.name]
         message = messages[-1]
         if "function_call" in message:
             _, func_return = self.execute_function(message["function_call"])
             return func_return
         if self._code_execution_config is False:
-            return default_reply if self.oai_config is False else self._oai_reply(messages)
+            return default_reply if self.llm_config is False else self._oai_reply(messages)
         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`
-            return default_reply if self.oai_config is False else self._oai_reply(messages)
+            return default_reply if self.llm_config is False else self._oai_reply(messages)
         # try to execute the code
         exitcode, logs = self.execute_code_blocks(code_blocks)
         exitcode2str = "execution succeeded" if exitcode == 0 else "execution failed"
@@ -387,7 +410,7 @@ def run_code(self, code, **kwargs):
         Returns:
             A tuple of (exitcode, logs, image).
             exitcode (int): the exit code of the code execution.
-            logs (bytes): the logs of the code execution.
+            logs (str): the logs of the code execution.
             image (str or None): the docker image used for the code execution.
         """
         return execute_code(code, **kwargs)
@@ -401,7 +424,6 @@ def execute_code_blocks(self, code_blocks):
                 lang = infer_lang(code)
             if lang in ["bash", "shell", "sh"]:
                 exitcode, logs, image = self.run_code(code, lang=lang, **self._code_execution_config)
-                logs = logs.decode("utf-8")
             elif lang in ["python", "Python"]:
                 if code.startswith("# filename: "):
                     filename = code[11 : code.find("\n")].strip()
@@ -412,7 +434,6 @@ def execute_code_blocks(self, code_blocks):
                     filename=filename,
                     **self._code_execution_config,
                 )
-                logs = logs.decode("utf-8")
             else:
                 # In case the language is not supported, we return an error message.
                 exitcode, logs, image = 1, f"unknown language {lang}", self._code_execution_config["use_docker"]
diff --git a/flaml/autogen/agentchat/user_proxy_agent.py b/flaml/autogen/agentchat/user_proxy_agent.py
index 76e8e61577..bf553ca5a7 100644
--- a/flaml/autogen/agentchat/user_proxy_agent.py
+++ b/flaml/autogen/agentchat/user_proxy_agent.py
@@ -6,7 +6,7 @@ class UserProxyAgent(ResponsiveAgent):
     """(Experimental) A proxy agent for the user, that can execute code and provide feedback to the other agents.
 
     UserProxyAgent is a subclass of ResponsiveAgent configured with `human_input_mode` to ALWAYS
-    and `oai_config` to False. By default, the agent will prompt for human input every time a message is received.
+    and `llm_config` to False. By default, the agent will prompt for human input every time a message is received.
     Code execution is enabled by default. LLM-based auto reply is disabled by default.
     To modify auto reply, override `generate_reply` method.
     To modify the way to get human input, override `get_human_input` method.
@@ -23,7 +23,7 @@ def __init__(
         human_input_mode: Optional[str] = "ALWAYS",
         function_map: Optional[Dict[str, Callable]] = None,
         code_execution_config: Optional[Union[Dict, bool]] = None,
-        oai_config: Optional[Union[Dict, bool]] = False,
+        llm_config: Optional[Union[Dict, bool]] = False,
         system_message: Optional[str] = "",
     ):
         """
@@ -59,12 +59,12 @@ def __init__(
                     If the code is executed in the current environment,
                     the code must be trusted.
                 - timeout (Optional, int): The maximum execution time in seconds.
-            oai_config (dict or False): oai inference configuration.
-                Please refer to [oai.Completion.create](/docs/reference/autogen/oai/completion#create)
+            llm_config (dict or False): llm inference configuration.
+                Please refer to [autogen.Completion.create](/docs/reference/autogen/oai/completion#create)
                 for available options.
-                Default to false, which disables oai-based auto reply.
-            system_message (str): system message for oai inference.
-                Only used when oai_config is not False. Use it to reprogram the agent.
+                Default to false, which disables llm-based auto reply.
+            system_message (str): system message for ChatCompletion inference.
+                Only used when llm_config is not False. Use it to reprogram the agent.
         """
         super().__init__(
             name,
@@ -74,5 +74,5 @@ def __init__(
             human_input_mode,
             function_map,
             code_execution_config,
-            oai_config,
+            llm_config,
         )
diff --git a/flaml/autogen/code_utils.py b/flaml/autogen/code_utils.py
index f721e93e0b..8cac7a455b 100644
--- a/flaml/autogen/code_utils.py
+++ b/flaml/autogen/code_utils.py
@@ -8,13 +8,20 @@
 import time
 from hashlib import md5
 import logging
-from flaml.autogen import oai, DEFAULT_MODEL, FAST_MODEL
+from flaml.autogen import oai
 
+try:
+    import docker
+except ImportError:
+    docker = None
+
+DEFAULT_MODEL = "gpt-4"
+FAST_MODEL = "gpt-3.5-turbo"
 # Regular expression for finding a code block
 CODE_BLOCK_PATTERN = r"```(\w*)\n(.*?)\n```"
 WORKING_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extensions")
 UNKNOWN = "unknown"
-TIMEOUT_MSG = bytes("Timeout", "utf-8")
+TIMEOUT_MSG = "Timeout"
 DEFAULT_TIMEOUT = 600
 
 
@@ -138,9 +145,9 @@ def execute_code(
     timeout: Optional[int] = None,
     filename: Optional[str] = None,
     work_dir: Optional[str] = None,
-    use_docker: Optional[Union[List[str], str, bool]] = True,
+    use_docker: Optional[Union[List[str], str, bool]] = docker is not None,
     lang: Optional[str] = "python",
-) -> Tuple[int, bytes, str]:
+) -> Tuple[int, str, str]:
     """Execute code in a docker container.
     This function is not tested on MacOS.
 
@@ -169,7 +176,7 @@ def execute_code(
 
     Returns:
         int: 0 if the code executes successfully.
-        bytes: The error message if the code fails to execute; the stdout otherwise.
+        str: The error message if the code fails to execute; the stdout otherwise.
         image: The docker image name after container run when docker is used.
     """
     assert code is not None or filename is not None, "Either code or filename must be provided."
@@ -216,9 +223,15 @@ def execute_code(
                 return 1, TIMEOUT_MSG, None
         if original_filename is None:
             os.remove(filepath)
-        return result.returncode, result.stderr if result.returncode else result.stdout, None
-
-    import docker
+            abs_path = str(pathlib.Path(filepath).absolute())
+        else:
+            abs_path = str(pathlib.Path(work_dir).absolute()) + "/"
+        if result.returncode:
+            logs = result.stderr.decode("utf-8")
+            logs = logs.replace(str(abs_path), "")
+        else:
+            logs = result.stdout.decode("utf-8")
+        return result.returncode, logs, None
 
     # create a docker client
     client = docker.from_env()
@@ -283,7 +296,8 @@ def execute_code(
     # get the container logs
     logs = container.logs().decode("utf-8").rstrip()
     # commit the image
-    container.commit(repository="python", tag=filename.replace("/", ""))
+    tag = filename.replace("/", "")
+    container.commit(repository="python", tag=tag)
     # remove the container
     container.remove()
     # check if the code executed successfully
@@ -296,11 +310,12 @@ def execute_code(
         # remove the exit code from the logs
         logs = pattern.sub("", logs)
 
-    logs = bytes(logs, "utf-8")
     if original_filename is None:
         os.remove(filepath)
+    if exit_code:
+        logs = logs.replace(f"/workspace/{filename if original_filename is None else ''}", "")
     # return the exit code, logs and image
-    return exit_code, logs, f"python:{filename}"
+    return exit_code, logs, f"python:{tag}"
 
 
 _GENERATE_ASSERTIONS_CONFIG = {
diff --git a/flaml/autogen/math_utils.py b/flaml/autogen/math_utils.py
index dc6bcde42c..7fcb42287f 100644
--- a/flaml/autogen/math_utils.py
+++ b/flaml/autogen/math_utils.py
@@ -9,7 +9,7 @@
 
 
 def solve_problem(problem: str, **config) -> str:
-    """(work in progress) Solve the math problem.
+    """(Experimental) Solve the math problem.
 
     Args:
         problem (str): The problem statement.
diff --git a/flaml/autogen/oai/completion.py b/flaml/autogen/oai/completion.py
index 200ecdcaf1..0c16705828 100644
--- a/flaml/autogen/oai/completion.py
+++ b/flaml/autogen/oai/completion.py
@@ -26,7 +26,7 @@
 
     ERROR = None
 except ImportError:
-    ERROR = ImportError("please install flaml[openai] option to use the flaml.oai subpackage.")
+    ERROR = ImportError("please install flaml[openai] option to use the flaml.autogen.oai subpackage.")
     openai_Completion = object
 logger = logging.getLogger(__name__)
 if not logger.handlers:
diff --git a/notebook/autogen_agentchat_MathChat.ipynb b/notebook/autogen_agentchat_MathChat.ipynb
index 8550e67ab7..d6dc1b183d 100644
--- a/notebook/autogen_agentchat_MathChat.ipynb
+++ b/notebook/autogen_agentchat_MathChat.ipynb
@@ -50,9 +50,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml import oai\n",
+    "from flaml import autogen\n",
     "\n",
-    "config_list = oai.config_list_from_json(\n",
+    "config_list = autogen.config_list_from_json(\n",
     "    \"OAI_CONFIG_LIST\",\n",
     "    filter_dict={\n",
     "        \"model\": {\n",
@@ -122,16 +122,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml.autogen.agentchat import AssistantAgent\n",
     "from flaml.autogen.agentchat.contrib.math_user_proxy_agent import MathUserProxyAgent\n",
     "\n",
-    "oai.ChatCompletion.start_logging()\n",
+    "autogen.ChatCompletion.start_logging()\n",
     "\n",
     "# 1. create an AssistantAgent instance named \"assistant\"\n",
-    "assistant = AssistantAgent(\n",
+    "assistant = autogen.AssistantAgent(\n",
     "    name=\"assistant\", \n",
     "    system_message=\"You are a helpful assistant.\",\n",
-    "    oai_config={\n",
+    "    llm_config={\n",
     "        \"request_timeout\": 600,\n",
     "        \"seed\": 42,\n",
     "        \"config_list\": config_list,\n",
@@ -856,7 +855,7 @@
     }
    ],
    "source": [
-    "assistant.reset() # clear LLM assistant's message history\n",
+    "assistant.reset()  # clear LLM assistant's message history\n",
     "\n",
     "# The wolfram alpha appid is required for this example (the assistant may choose to query Wolfram Alpha).\n",
     "import os\n",
diff --git a/notebook/autogen_agentchat_auto_feedback_from_code_execution.ipynb b/notebook/autogen_agentchat_auto_feedback_from_code_execution.ipynb
index 4c9ef96ce0..8a5d914868 100644
--- a/notebook/autogen_agentchat_auto_feedback_from_code_execution.ipynb
+++ b/notebook/autogen_agentchat_auto_feedback_from_code_execution.ipynb
@@ -64,9 +64,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml import oai\n",
+    "from flaml import autogen\n",
     "\n",
-    "config_list = oai.config_list_from_json(\n",
+    "config_list = autogen.config_list_from_json(\n",
     "    \"OAI_CONFIG_LIST\",\n",
     "    filter_dict={\n",
     "        \"model\": [\"gpt-4\", \"gpt4\", \"gpt-4-32k\", \"gpt-4-32k-0314\"],\n",
@@ -192,19 +192,17 @@
     }
    ],
    "source": [
-    "from flaml.autogen.agentchat import AssistantAgent, UserProxyAgent\n",
-    "\n",
     "# create an AssistantAgent named \"assistant\"\n",
-    "assistant = AssistantAgent(\n",
+    "assistant = autogen.AssistantAgent(\n",
     "    name=\"assistant\",\n",
-    "    oai_config={\n",
+    "    llm_config={\n",
     "        \"seed\": 42,  # seed for caching and reproducibility\n",
     "        \"config_list\": config_list,  # a list of OpenAI API configurations\n",
     "        \"temperature\": 0,  # temperature for sampling\n",
-    "    },  # configuration for flaml.oai, an enhanced inference API compatible with OpenAI API\n",
+    "    },  # configuration for autogen's enhanced inference API which is compatible with OpenAI API\n",
     ")\n",
     "# create a UserProxyAgent instance named \"user_proxy\"\n",
-    "user_proxy = UserProxyAgent(\n",
+    "user_proxy = autogen.UserProxyAgent(\n",
     "    name=\"user_proxy\",\n",
     "    human_input_mode=\"NEVER\",\n",
     "    max_consecutive_auto_reply=10,\n",
@@ -348,6 +346,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -367,7 +366,7 @@
     "from typing import Dict, Union\n",
     "from IPython import get_ipython\n",
     "\n",
-    "class IPythonUserProxyAgent(UserProxyAgent):\n",
+    "class IPythonUserProxyAgent(autogen.UserProxyAgent):\n",
     "    def __init__(self, name: str, **kwargs):\n",
     "        super().__init__(name, **kwargs)\n",
     "        self._ipython = get_ipython()\n",
@@ -386,10 +385,11 @@
     "        if result.error_in_exec is not None:\n",
     "            log += f\"\\n{result.error_in_exec}\"\n",
     "            exitcode = 1\n",
-    "        return exitcode, bytes(log, \"utf-8\"), None"
+    "        return exitcode, log, None"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
diff --git a/notebook/autogen_agentchat_function_call.ipynb b/notebook/autogen_agentchat_function_call.ipynb
index 964879af8e..abf786ddad 100644
--- a/notebook/autogen_agentchat_function_call.ipynb
+++ b/notebook/autogen_agentchat_function_call.ipynb
@@ -66,9 +66,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml import oai\n",
+    "from flaml import autogen\n",
     "\n",
-    "config_list = oai.config_list_from_models(model_list=[\"gpt-4\", \"gpt-3.5-turbo\", \"gpt-3.5-turbo-16k\"], exclude=\"aoai\")"
+    "config_list = autogen.config_list_from_models(model_list=[\"gpt-4\", \"gpt-3.5-turbo\", \"gpt-3.5-turbo-16k\"], exclude=\"aoai\")"
    ]
   },
   {
@@ -117,14 +117,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "user_proxy (to chatbot):\n",
+      "\u001b[33muser_proxy\u001b[0m (to chatbot):\n",
       "\n",
       "Draw two agents chatting with each other with an example dialog. Don't add plt.show().\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "chatbot (to user_proxy):\n",
+      "\u001b[33mchatbot\u001b[0m (to user_proxy):\n",
       "\n",
-      "***** Suggested function Call: python *****\n",
+      "\u001b[32m***** Suggested function Call: python *****\u001b[0m\n",
       "Arguments: \n",
       "{\n",
       "  \"cell\": \"\n",
@@ -156,25 +156,25 @@
       "plt.axis('off')\n",
       "\"\n",
       "}\n",
-      "*******************************************\n",
+      "\u001b[32m*******************************************\u001b[0m\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "user_proxy (to chatbot):\n",
+      "\u001b[33muser_proxy\u001b[0m (to chatbot):\n",
       "\n",
-      "***** Response from calling function \"python\" *****\n",
+      "\u001b[32m***** Response from calling function \"python\" *****\u001b[0m\n",
       "Error: Invalid \\escape: line 1 column 785 (char 784)\n",
       " You argument should follow json format.\n",
-      "***************************************************\n",
+      "\u001b[32m***************************************************\u001b[0m\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "chatbot (to user_proxy):\n",
+      "\u001b[33mchatbot\u001b[0m (to user_proxy):\n",
       "\n",
-      "***** Suggested function Call: python *****\n",
+      "\u001b[32m***** Suggested function Call: python *****\u001b[0m\n",
       "Arguments: \n",
       "{\n",
       "  \"cell\": \"import matplotlib.pyplot as plt\\nimport matplotlib.patches as mpatches\\n\\n# Define basic parameters\\nface_color = '#FFDDC1'\\nplt.figure(figsize=(10, 2))\\n\\n# Agent 1\\nagent1 = mpatches.FancyBboxPatch((0.02, 0.4), 0.2, 0.6, boxstyle=mpatches.BoxStyle('Round', pad=0.02))\\nplt.gca().add_artist(agent1)\\nplt.gca().text(0.12, 0.7, 'Agent 1', ha='center', va='center', fontsize=12, color='blue')\\n\\n# Agent 2\\nagent2 = mpatches.FancyBboxPatch((0.45, 0.4), 0.2, 0.6, boxstyle=mpatches.BoxStyle('Round', pad=0.02))\\nplt.gca().add_artist(agent2)\\nplt.gca().text(0.55, 0.7, 'Agent 2', ha='center', va='center', fontsize=12, color='red')\\n\\n# Dialog\\nplt.gca().text(0.12, 0.35, '\\\"Hello, how are you?\\\"', ha='center', va='center', fontsize=10)\\nplt.gca().text(0.55, 0.15, '\\\"I\\\\'m fine, thank you!\\\"', ha='center', va='center', fontsize=10)\\n\\n# Descriptions\\nplt.gca().text(0.12, 0.15, 'Greeting', ha='center', va='center', fontsize=10)\\nplt.gca().text(0.55, 0.35, 'Response', ha='center', va='center', fontsize=10)\\n\\nplt.axis('off')\"\n",
       "}\n",
-      "*******************************************\n",
+      "\u001b[32m*******************************************\u001b[0m\n",
       "\n",
       "--------------------------------------------------------------------------------\n"
      ]
@@ -203,14 +203,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "user_proxy (to chatbot):\n",
+      "\u001b[33muser_proxy\u001b[0m (to chatbot):\n",
       "\n",
-      "***** Response from calling function \"python\" *****\n",
+      "\u001b[32m***** Response from calling function \"python\" *****\u001b[0m\n",
       "(0.0, 1.0, 0.0, 1.0)\n",
-      "***************************************************\n",
+      "\u001b[32m***************************************************\u001b[0m\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "chatbot (to user_proxy):\n",
+      "\u001b[33mchatbot\u001b[0m (to user_proxy):\n",
       "\n",
       "TERMINATE\n",
       "\n",
@@ -219,9 +219,7 @@
     }
    ],
    "source": [
-    "from flaml.autogen.agentchat import AssistantAgent, UserProxyAgent\n",
-    "\n",
-    "oai_config = {\n",
+    "llm_config = {\n",
     "    \"functions\": [\n",
     "        {\n",
     "            \"name\": \"python\",\n",
@@ -255,14 +253,14 @@
     "    \"config_list\": config_list,\n",
     "    \"request_timeout\": 120,\n",
     "}\n",
-    "chatbot = AssistantAgent(\n",
+    "chatbot = autogen.AssistantAgent(\n",
     "    name=\"chatbot\",\n",
     "    system_message=\"For coding tasks, only use the functions you have been provided with. Reply TERMINATE when the task is done.\",\n",
-    "    oai_config=oai_config,\n",
+    "    llm_config=llm_config,\n",
     ")\n",
     "\n",
     "# create a UserProxyAgent instance named \"user_proxy\"\n",
-    "user_proxy = UserProxyAgent(\n",
+    "user_proxy = autogen.UserProxyAgent(\n",
     "    name=\"user_proxy\",\n",
     "    is_termination_msg=lambda x: x.get(\"content\", \"\") and x.get(\"content\", \"\").rstrip().endswith(\"TERMINATE\"),\n",
     "    human_input_mode=\"NEVER\",\n",
@@ -314,7 +312,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "id": "4a917492",
    "metadata": {},
    "outputs": [
@@ -322,19 +320,19 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "user_proxy (to chatbot):\n",
+      "\u001b[33muser_proxy\u001b[0m (to chatbot):\n",
       "\n",
       "Problem: Find all $x$ that satisfy the inequality $(2x+10)(x+3)<(3x+9)(x+8)$. Express your answer in interval notation.\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "chatbot (to user_proxy):\n",
+      "\u001b[33mchatbot\u001b[0m (to user_proxy):\n",
       "\n",
-      "***** Suggested function Call: query_wolfram *****\n",
+      "\u001b[32m***** Suggested function Call: query_wolfram *****\u001b[0m\n",
       "Arguments: \n",
       "{\n",
       "  \"query\": \"solve (2x+10)(x+3)<(3x+9)(x+8) for x\"\n",
       "}\n",
-      "**************************************************\n",
+      "\u001b[32m**************************************************\u001b[0m\n",
       "\n",
       "--------------------------------------------------------------------------------\n"
      ]
@@ -343,24 +341,24 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "user_proxy (to chatbot):\n",
+      "\u001b[33muser_proxy\u001b[0m (to chatbot):\n",
       "\n",
-      "***** Response from calling function \"query_wolfram\" *****\n",
+      "\u001b[32m***** Response from calling function \"query_wolfram\" *****\u001b[0m\n",
       "('Assumption: solve (2 x + 10) (x + 3)<(3 x + 9) (x + 8) for x \\nAnswer: ans 0: x<-14\\nans 1: x>-3\\n', True)\n",
-      "**********************************************************\n",
+      "\u001b[32m**********************************************************\u001b[0m\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "chatbot (to user_proxy):\n",
+      "\u001b[33mchatbot\u001b[0m (to user_proxy):\n",
       "\n",
       "The solution to the inequality $(2x+10)(x+3)<(3x+9)(x+8)$ is $x \\in (-\\infty, -14) \\cup (-3, +\\infty)$. TERMINATE\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "user_proxy (to chatbot):\n",
+      "\u001b[33muser_proxy\u001b[0m (to chatbot):\n",
       "\n",
       "\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "chatbot (to user_proxy):\n",
+      "\u001b[33mchatbot\u001b[0m (to user_proxy):\n",
       "\n",
       "TERMINATE\n",
       "\n",
@@ -370,14 +368,13 @@
    ],
    "source": [
     "import os\n",
-    "from flaml.autogen.agentchat import AssistantAgent, UserProxyAgent\n",
     "from flaml.autogen.agentchat.contrib.math_user_proxy_agent import MathUserProxyAgent\n",
     "\n",
     "# you need to provide a wolfram alpha appid to run this example\n",
     "if not os.environ.get(\"WOLFRAM_ALPHA_APPID\"):\n",
     "    os.environ[\"WOLFRAM_ALPHA_APPID\"] = open(\"wolfram.txt\").read().strip()\n",
     "\n",
-    "oai_config = {\n",
+    "llm_config = {\n",
     "    \"model\": \"gpt-4-0613\",\n",
     "    \"functions\": [\n",
     "        {\n",
@@ -397,15 +394,15 @@
     "    ],\n",
     "    \"config_list\": config_list,\n",
     "}\n",
-    "chatbot = AssistantAgent(\n",
+    "chatbot = autogen.AssistantAgent(\n",
     "    name=\"chatbot\",\n",
     "    system_message=\"Only use the functions you have been provided with. Do not ask user to perform other actions than executing the functions. Reply TERMINATE when the task is done.\",\n",
-    "    oai_config=oai_config,\n",
+    "    llm_config=llm_config,\n",
     ")\n",
     "\n",
     "# the key in `function_map` should match the function name in \"functions\" above\n",
     "# we register a class instance method directly\n",
-    "user_proxy = UserProxyAgent(\n",
+    "user_proxy = autogen.UserProxyAgent(\n",
     "    \"user_proxy\",\n",
     "    max_consecutive_auto_reply=2,\n",
     "    human_input_mode=\"NEVER\",\n",
diff --git a/notebook/autogen_agentchat_human_feedback.ipynb b/notebook/autogen_agentchat_human_feedback.ipynb
index 1bf1fc8046..3d2a4e0122 100644
--- a/notebook/autogen_agentchat_human_feedback.ipynb
+++ b/notebook/autogen_agentchat_human_feedback.ipynb
@@ -64,9 +64,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml import oai\n",
+    "from flaml import autogen\n",
     "\n",
-    "config_list = oai.config_list_from_json(\"OAI_CONFIG_LIST\")"
+    "config_list = autogen.config_list_from_json(\"OAI_CONFIG_LIST\")"
    ]
   },
   {
@@ -139,25 +139,23 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml.autogen.agentchat import AssistantAgent, UserProxyAgent\n",
-    "\n",
     "# create an AssistantAgent instance named \"assistant\"\n",
-    "assistant = AssistantAgent(\n",
+    "assistant = autogen.AssistantAgent(\n",
     "    name=\"assistant\",\n",
-    "    oai_config={\n",
+    "    llm_config={\n",
     "        \"seed\": 41,\n",
     "        \"config_list\": config_list,\n",
     "    }\n",
     ")\n",
     "# create a UserProxyAgent instance named \"user_proxy\"\n",
-    "user_proxy = UserProxyAgent(\n",
+    "user_proxy = autogen.UserProxyAgent(\n",
     "    name=\"user_proxy\",\n",
     "    human_input_mode=\"ALWAYS\",\n",
     "    is_termination_msg=lambda x: x.get(\"content\", \"\").rstrip().endswith(\"TERMINATE\"),\n",
     ")\n",
     "\n",
     "# the purpose of the following line is to log the conversation history\n",
-    "oai.ChatCompletion.start_logging()\n"
+    "autogen.ChatCompletion.start_logging()\n"
    ]
   },
   {
@@ -380,7 +378,7 @@
     "\n",
     "The human user can provide feedback at each step. When the human user didn't provide feedback, the code was executed. The executed results and error messages are returned to the assistant, and the assistant was able to modify the code based on the feedback. In the end, the task is complete and a \"TERMINATE\" signal is sent from the assistant. The user skipped feedback in the end and the conversation is finished.\n",
     "\n",
-    "After the conversation is finished, we can save the log of the conversation between the two agents. The log can be accessed from `oai.ChatCompletion.logged_history`."
+    "After the conversation is finished, we can save the log of the conversation between the two agents. The log can be accessed from `autogen.ChatCompletion.logged_history`."
    ]
   },
   {
@@ -397,7 +395,7 @@
     }
    ],
    "source": [
-    "print(oai.ChatCompletion.logged_history)"
+    "print(autogen.ChatCompletion.logged_history)"
    ]
   },
   {
@@ -408,7 +406,7 @@
    "source": [
     "import json\n",
     "\n",
-    "json.dump(oai.ChatCompletion.logged_history, open(\"conversations.json\", \"w\"), indent=2)"
+    "json.dump(autogen.ChatCompletion.logged_history, open(\"conversations.json\", \"w\"), indent=2)"
    ]
   }
  ],
diff --git a/notebook/autogen_optiguide.ipynb b/notebook/autogen_agentchat_optiguide.ipynb
similarity index 99%
rename from notebook/autogen_optiguide.ipynb
rename to notebook/autogen_agentchat_optiguide.ipynb
index fd1f5c7478..ceb128a26b 100644
--- a/notebook/autogen_optiguide.ipynb
+++ b/notebook/autogen_agentchat_optiguide.ipynb
@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "a461c72d",
    "metadata": {},
@@ -11,6 +12,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "59a05fc7",
    "metadata": {},
@@ -22,6 +24,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "5e92d200",
    "metadata": {},
@@ -37,6 +40,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "18c316b5",
    "metadata": {},
@@ -65,6 +69,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "8b7f90c8",
    "metadata": {},
@@ -147,6 +152,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "e9e7e728",
    "metadata": {},
@@ -235,6 +241,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "5a5a7d7e",
    "metadata": {},
@@ -616,6 +623,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "3489022b",
    "metadata": {},
diff --git a/notebook/autogen_agentchat_planning.ipynb b/notebook/autogen_agentchat_planning.ipynb
index b50d33a217..be81ee491a 100644
--- a/notebook/autogen_agentchat_planning.ipynb
+++ b/notebook/autogen_agentchat_planning.ipynb
@@ -25,9 +25,9 @@
     "\n",
     "## Requirements\n",
     "\n",
-    "FLAML requires `Python>=3.8`. To run this notebook example, please install flaml with the [autogen] option:\n",
+    "FLAML requires `Python>=3.8`. To run this notebook example, please install flaml with the [autogen] option and docker:\n",
     "```bash\n",
-    "pip install flaml[autogen]\n",
+    "pip install flaml[autogen] docker\n",
     "```"
    ]
   },
@@ -44,7 +44,7 @@
    },
    "outputs": [],
    "source": [
-    "# %pip install flaml[autogen]~=2.0.0rc4"
+    "# %pip install flaml[autogen]~=2.0.0rc4 docker"
    ]
   },
   {
@@ -69,9 +69,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml import oai\n",
+    "from flaml import autogen\n",
     "\n",
-    "config_list = oai.config_list_from_json(\n",
+    "config_list = autogen.config_list_from_json(\n",
     "    \"OAI_CONFIG_LIST\",\n",
     "    filter_dict={\n",
     "        \"model\": [\"gpt-4\", \"gpt4\", \"gpt-4-32k\", \"gpt-4-32k-0314\"],\n",
@@ -123,15 +123,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml.autogen.agentchat import AssistantAgent, UserProxyAgent\n",
-    "\n",
-    "planner = AssistantAgent(\n",
+    "planner = autogen.AssistantAgent(\n",
     "    name=\"planner\",\n",
-    "    oai_config={\"config_list\": config_list},\n",
+    "    llm_config={\"config_list\": config_list},\n",
     "    # the default system message of the AssistantAgent is overwritten here\n",
     "    system_message=\"You are a helpful AI assistant. You suggest coding and reasoning steps for another AI assistant to accomplish a task. Do not suggest concrete code. For any action beyond writing code or reasoning, convert it to a step which can be implemented by writing code. For example, the action of browsing the web can be implemented by writing code which reads and prints the content of a web page. Finally, inspect the execution result. If the plan is not good, suggest a better plan. If the execution is wrong, analyze the error and suggest a fix.\"\n",
     ")\n",
-    "planner_user = UserProxyAgent(\n",
+    "planner_user = autogen.UserProxyAgent(\n",
     "    name=\"planner_user\",\n",
     "    max_consecutive_auto_reply=0,  # terminate without auto-reply\n",
     "    human_input_mode=\"NEVER\",\n",
@@ -158,14 +156,14 @@
    "outputs": [],
    "source": [
     "# create an AssistantAgent instance named \"assistant\"\n",
-    "assistant = AssistantAgent(\n",
+    "assistant = autogen.AssistantAgent(\n",
     "    name=\"assistant\",\n",
-    "    oai_config={\n",
+    "    llm_config={\n",
     "        \"temperature\": 0,\n",
     "        \"request_timeout\": 600,\n",
     "        \"seed\": 42,\n",
     "        \"model\": \"gpt-4-0613\",\n",
-    "        \"config_list\": oai.config_list_openai_aoai(exclude=\"aoai\"),\n",
+    "        \"config_list\": autogen.config_list_openai_aoai(exclude=\"aoai\"),\n",
     "        \"functions\": [\n",
     "            {\n",
     "                \"name\": \"ask_planner\",\n",
@@ -186,7 +184,7 @@
     ")\n",
     "\n",
     "# create a UserProxyAgent instance named \"user_proxy\"\n",
-    "user_proxy = UserProxyAgent(\n",
+    "user_proxy = autogen.UserProxyAgent(\n",
     "    name=\"user_proxy\",\n",
     "    human_input_mode=\"TERMINATE\",\n",
     "    max_consecutive_auto_reply=10,\n",
@@ -215,12 +213,12 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "user_proxy (to assistant):\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
       "Suggest a fix to an open good first issue of flaml\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
       "To suggest a fix to an open good first issue of FLAML, we first need to fetch the list of open issues labeled as \"good first issue\" from the FLAML GitHub repository. We can do this using the GitHub API. Here is a Python script that fetches and prints the list of such issues.\n",
       "\n",
@@ -244,15 +242,15 @@
       "Please save this script as `fetch_issues.py` and run it. It will print the list of open issues labeled as \"good first issue\" in the FLAML repository.\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> USING AUTO REPLY...\n"
+      "\u001b[31m\n",
+      ">>>>>>>> USING AUTO REPLY...\u001b[0m\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "user_proxy (to assistant):\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
       "exitcode: 0 (execution succeeded)\n",
       "Code output: \n",
@@ -278,7 +276,7 @@
       "\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
       "***** Suggested function Call: ask_planner *****\n",
       "Arguments: \n",
@@ -288,9 +286,9 @@
       "************************************************\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> USING AUTO REPLY...\n",
-      "planner_user (to planner):\n",
+      "\u001b[31m\n",
+      ">>>>>>>> USING AUTO REPLY...\u001b[0m\n",
+      "\u001b[33mplanner_user\u001b[0m (to planner):\n",
       "\n",
       "Here are the open good first issues of FLAML: \n",
       "\n",
@@ -317,7 +315,7 @@
       "Please suggest a fix for one of these issues.\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "planner (to planner_user):\n",
+      "\u001b[33mplanner\u001b[0m (to planner_user):\n",
       "\n",
       "Without certain details about the issues, it's a bit challenging to suggest explicit fixes. However, I'll give you a roadmap to tackle some of these:\n",
       "\n",
@@ -354,7 +352,7 @@
       "Remember to thoroughly inspect the execution result for each issue to ensure your fix is effective. If the plan does not appear to fix the issue, adjust your plan based on the error message or unexpected behavior.\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "user_proxy (to assistant):\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
       "***** Response from calling function \"ask_planner\" *****\n",
       "Without certain details about the issues, it's a bit challenging to suggest explicit fixes. However, I'll give you a roadmap to tackle some of these:\n",
@@ -393,7 +391,7 @@
       "********************************************************\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
       "Based on the planner's suggestions, let's tackle the first issue: \"Issue #1120: use_label_encoder warning with xgboost\". \n",
       "\n",
@@ -414,20 +412,20 @@
       "TERMINATE\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> USING AUTO REPLY...\n",
-      "user_proxy (to assistant):\n",
+      "\u001b[31m\n",
+      ">>>>>>>> USING AUTO REPLY...\u001b[0m\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
       "\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
       "TERMINATE\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> NO HUMAN INPUT RECEIVED.\n"
+      "\u001b[31m\n",
+      ">>>>>>>> NO HUMAN INPUT RECEIVED.\u001b[0m\n"
      ]
     }
    ],
diff --git a/notebook/autogen_agentchat_two_users.ipynb b/notebook/autogen_agentchat_two_users.ipynb
index 781ddc124f..92f63d82ec 100644
--- a/notebook/autogen_agentchat_two_users.ipynb
+++ b/notebook/autogen_agentchat_two_users.ipynb
@@ -69,9 +69,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml import oai\n",
+    "from flaml import autogen\n",
     "\n",
-    "config_list = oai.config_list_from_json(\n",
+    "config_list = autogen.config_list_from_json(\n",
     "    \"OAI_CONFIG_LIST\",\n",
     "    filter_dict={\n",
     "        \"model\": [\"gpt-4\", \"gpt4\", \"gpt-4-32k\", \"gpt-4-32k-0314\"],\n",
@@ -123,17 +123,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml.autogen.agentchat import AssistantAgent, UserProxyAgent\n",
-    "\n",
     "def ask_expert(message):\n",
-    "    assistant_for_expert = AssistantAgent(\n",
+    "    assistant_for_expert = autogen.AssistantAgent(\n",
     "        name=\"assistant_for_expert\",\n",
-    "        oai_config={\n",
+    "        llm_config={\n",
     "            \"temperature\": 0,\n",
     "            \"config_list\": config_list,\n",
     "        },\n",
     "    )\n",
-    "    expert = UserProxyAgent(\n",
+    "    expert = autogen.UserProxyAgent(\n",
     "        name=\"expert\",\n",
     "        human_input_mode=\"ALWAYS\",\n",
     "        code_execution_config={\"work_dir\": \"expert\"},\n",
@@ -163,14 +161,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "assistant_for_student = AssistantAgent(\n",
+    "assistant_for_student = autogen.AssistantAgent(\n",
     "    name=\"assistant_for_student\",\n",
-    "    oai_config={\n",
+    "    llm_config={\n",
     "        \"request_timeout\": 600,\n",
     "        \"seed\": 42,\n",
     "        # Excluding azure openai endpoints from the config list.\n",
     "        # Change to `exclude=\"openai\"` to exclude openai endpoints, or remove the `exclude` argument to include both.\n",
-    "        \"config_list\": oai.config_list_openai_aoai(exclude=\"aoai\"),\n",
+    "        \"config_list\": autogen.config_list_openai_aoai(exclude=\"aoai\"),\n",
     "        \"model\": \"gpt-4-0613\",  # make sure the endpoint you use supports the model\n",
     "        \"temperature\": 0,\n",
     "        \"functions\": [\n",
@@ -192,7 +190,7 @@
     "    }\n",
     ")\n",
     "\n",
-    "student = UserProxyAgent(\n",
+    "student = autogen.UserProxyAgent(\n",
     "    name=\"student\",\n",
     "    human_input_mode=\"TERMINATE\",\n",
     "    max_consecutive_auto_reply=10,\n",
diff --git a/notebook/autogen_agentchat_web_info.ipynb b/notebook/autogen_agentchat_web_info.ipynb
index 079d507f43..75431c6571 100644
--- a/notebook/autogen_agentchat_web_info.ipynb
+++ b/notebook/autogen_agentchat_web_info.ipynb
@@ -30,7 +30,7 @@
     "\n",
     "## Requirements\n",
     "\n",
-    "FLAML requires `Python>=3.8`. To run this notebook example, please install flaml with the [autogen] option:\n",
+    "FLAML requires `Python>=3.8`. To run this notebook example, please install flaml with the [autogen] option and docker:\n",
     "```bash\n",
     "pip install flaml[autogen]\n",
     "```"
@@ -49,7 +49,7 @@
    },
    "outputs": [],
    "source": [
-    "# %pip install flaml[autogen]~=2.0.0rc4"
+    "# %pip install flaml[autogen]~=2.0.0rc4 docker"
    ]
   },
   {
@@ -68,14 +68,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml import oai\n",
+    "from flaml import autogen\n",
     "\n",
-    "config_list = oai.config_list_from_json(\n",
+    "config_list = autogen.config_list_from_json(\n",
     "    \"OAI_CONFIG_LIST\",\n",
     "    filter_dict={\n",
     "        \"model\": [\"gpt4\", \"gpt-4-32k\", \"gpt-4-32k-0314\"],\n",
     "    },\n",
-    ")"
+    ")\n",
+    "\n",
+    "llm_config={\n",
+    "    \"request_timeout\": 600,\n",
+    "    \"seed\": 42,\n",
+    "    \"config_list\": config_list,\n",
+    "    \"model\": \"gpt-4-32k\",  # modify if the endpoint you use doesn't support this model\n",
+    "    \"temperature\": 0,\n",
+    "}"
    ]
   },
   {
@@ -130,26 +138,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flaml.autogen.agentchat import AssistantAgent, UserProxyAgent\n",
-    "\n",
     "# create an AssistantAgent instance named \"assistant\"\n",
-    "assistant = AssistantAgent(\n",
+    "assistant = autogen.AssistantAgent(\n",
     "    name=\"assistant\",\n",
-    "    oai_config={\n",
-    "        \"request_timeout\": 600,\n",
-    "        \"seed\": 42,\n",
-    "        \"config_list\": config_list,\n",
-    "        \"model\": \"gpt-4-32k\",  # modify if the endpoint you use doesn't support this model\n",
-    "        \"temperature\": 0,\n",
-    "    }\n",
+    "    llm_config=llm_config,\n",
     ")\n",
     "# create a UserProxyAgent instance named \"user_proxy\"\n",
-    "user_proxy = UserProxyAgent(\n",
+    "user_proxy = autogen.UserProxyAgent(\n",
     "    name=\"user_proxy\",\n",
     "    human_input_mode=\"TERMINATE\",\n",
     "    max_consecutive_auto_reply=10,\n",
     "    is_termination_msg=lambda x: x.get(\"content\", \"\").rstrip().endswith(\"TERMINATE\"),\n",
     "    code_execution_config={\"work_dir\": \"web\"},\n",
+    "    llm_config=llm_config,\n",
+    "    system_message=\"\"\"Reply TERMINATE if the task has been solved at full satisfaction.\n",
+    "Otherwise, reply CONTINUE, or the reason why the task is not solved yet.\"\"\"\n",
     ")"
    ]
   },
@@ -172,14 +175,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "user_proxy (to assistant):\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
       "\n",
       "Who should read this paper: https://arxiv.org/abs/2306.01337\n",
       "\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
       "To determine who should read this paper, let's first extract the abstract and analyze its content. We can use Python with the `requests` and `BeautifulSoup` libraries to do this. Please execute the following code:\n",
       "\n",
@@ -197,15 +200,15 @@
       "After executing the code, I will analyze the abstract and suggest who should read the paper.\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> USING AUTO REPLY FOR THE USER...\n"
+      "\u001b[31m\n",
+      ">>>>>>>> USING AUTO REPLY...\u001b[0m\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "user_proxy (to assistant):\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
       "exitcode: 0 (execution succeeded)\n",
       "Code output: \n",
@@ -223,7 +226,7 @@
       "\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
       "Based on the abstract, the paper is about using Large Language Models (LLMs), specifically GPT-4, to solve complex mathematical problems. The authors propose a conversational problem-solving framework called MathChat and evaluate it on difficult high school competition problems from the MATH dataset.\n",
       "\n",
@@ -236,7 +239,9 @@
       "\n",
       "If you belong to any of these categories or have a general interest in AI and its applications, you might find this paper valuable. TERMINATE\n",
       "\n",
-      "--------------------------------------------------------------------------------\n"
+      "--------------------------------------------------------------------------------\n",
+      "\u001b[31m\n",
+      ">>>>>>>> NO HUMAN INPUT RECEIVED.\u001b[0m\n"
      ]
     }
    ],
@@ -267,12 +272,12 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "user_proxy (to assistant):\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
       "Show me the YTD gain of 10 largest technology companies as of today.\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
       "To get the YTD gain of the 10 largest technology companies, we can use a Python script to fetch the data from a financial API like Yahoo Finance. Please execute the following Python code:\n",
       "\n",
@@ -310,9 +315,9 @@
       "This code will fetch the 10 largest technology companies from TradingView and then calculate their YTD gains using Yahoo Finance data.\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> USING AUTO REPLY FOR THE USER...\n",
-      "user_proxy (to assistant):\n",
+      "\u001b[31m\n",
+      ">>>>>>>> USING AUTO REPLY...\u001b[0m\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
       "exitcode: 1 (execution failed)\n",
       "Code output: \n",
@@ -325,7 +330,7 @@
       "\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
       "I apologize for the error. It seems that the website structure has changed, and the script is unable to find the table containing the top technology companies. Let's try an alternative approach using the Yahoo Finance library. Please install the library first by executing:\n",
       "\n",
@@ -360,32 +365,31 @@
       "This code uses a predefined list of the 10 largest technology companies and calculates their YTD gains using the Yahoo Finance library.\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> USING AUTO REPLY FOR THE USER...\n",
-      "user_proxy (to assistant):\n",
+      "\u001b[31m\n",
+      ">>>>>>>> USING AUTO REPLY...\u001b[0m\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
       "exitcode: 1 (execution failed)\n",
       "Code output: \n",
       "Defaulting to user installation because normal site-packages is not writeable\n",
       "Requirement already satisfied: yfinance in /home/vscode/.local/lib/python3.9/site-packages (0.2.26)\n",
-      "Requirement already satisfied: lxml>=4.9.1 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (4.9.3)\n",
-      "Requirement already satisfied: pytz>=2022.5 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (2023.3)\n",
       "Requirement already satisfied: numpy>=1.16.5 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (1.25.1)\n",
-      "Requirement already satisfied: html5lib>=1.1 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (1.1)\n",
+      "Requirement already satisfied: pandas>=1.3.0 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (1.3.3)\n",
       "Requirement already satisfied: beautifulsoup4>=4.11.1 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (4.12.2)\n",
-      "Requirement already satisfied: multitasking>=0.0.7 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (0.0.11)\n",
+      "Requirement already satisfied: html5lib>=1.1 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (1.1)\n",
+      "Requirement already satisfied: pytz>=2022.5 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (2023.3)\n",
       "Requirement already satisfied: appdirs>=1.4.4 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (1.4.4)\n",
-      "Requirement already satisfied: pandas>=1.3.0 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (2.0.3)\n",
+      "Requirement already satisfied: lxml>=4.9.1 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (4.9.3)\n",
       "Requirement already satisfied: requests>=2.31 in /usr/local/lib/python3.9/site-packages (from yfinance) (2.31.0)\n",
       "Requirement already satisfied: frozendict>=2.3.4 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (2.3.8)\n",
+      "Requirement already satisfied: multitasking>=0.0.7 in /home/vscode/.local/lib/python3.9/site-packages (from yfinance) (0.0.11)\n",
       "Requirement already satisfied: soupsieve>1.2 in /home/vscode/.local/lib/python3.9/site-packages (from beautifulsoup4>=4.11.1->yfinance) (2.4.1)\n",
       "Requirement already satisfied: six>=1.9 in /usr/local/lib/python3.9/site-packages (from html5lib>=1.1->yfinance) (1.16.0)\n",
       "Requirement already satisfied: webencodings in /home/vscode/.local/lib/python3.9/site-packages (from html5lib>=1.1->yfinance) (0.5.1)\n",
-      "Requirement already satisfied: python-dateutil>=2.8.2 in /home/vscode/.local/lib/python3.9/site-packages (from pandas>=1.3.0->yfinance) (2.8.2)\n",
-      "Requirement already satisfied: tzdata>=2022.1 in /home/vscode/.local/lib/python3.9/site-packages (from pandas>=1.3.0->yfinance) (2023.3)\n",
-      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.9/site-packages (from requests>=2.31->yfinance) (2.0.3)\n",
-      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.9/site-packages (from requests>=2.31->yfinance) (2023.5.7)\n",
+      "Requirement already satisfied: python-dateutil>=2.7.3 in /home/vscode/.local/lib/python3.9/site-packages (from pandas>=1.3.0->yfinance) (2.8.2)\n",
       "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.9/site-packages (from requests>=2.31->yfinance) (3.4)\n",
+      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.9/site-packages (from requests>=2.31->yfinance) (2023.5.7)\n",
+      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.9/site-packages (from requests>=2.31->yfinance) (2.0.3)\n",
       "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.9/site-packages (from requests>=2.31->yfinance) (3.2.0)\n",
       "\n",
       "FB: No timezone found, symbol may be delisted\n",
@@ -396,17 +400,17 @@
       "    ytd_gains = {symbol: get_ytd_gain(symbol) for symbol in top_tech_companies}\n",
       "  File \"/workspaces/FLAML/notebook/web/tmp_code_52c69e134e808cf8ba196ac60faede00.py\", line 13, in get_ytd_gain\n",
       "    ytd_gain = (data[\"Close\"].iloc[-1] - data[\"Close\"].iloc[0]) / data[\"Close\"].iloc[0] * 100\n",
-      "  File \"/home/vscode/.local/lib/python3.9/site-packages/pandas/core/indexing.py\", line 1103, in __getitem__\n",
+      "  File \"/home/vscode/.local/lib/python3.9/site-packages/pandas/core/indexing.py\", line 931, in __getitem__\n",
       "    return self._getitem_axis(maybe_callable, axis=axis)\n",
-      "  File \"/home/vscode/.local/lib/python3.9/site-packages/pandas/core/indexing.py\", line 1656, in _getitem_axis\n",
+      "  File \"/home/vscode/.local/lib/python3.9/site-packages/pandas/core/indexing.py\", line 1566, in _getitem_axis\n",
       "    self._validate_integer(key, axis)\n",
-      "  File \"/home/vscode/.local/lib/python3.9/site-packages/pandas/core/indexing.py\", line 1589, in _validate_integer\n",
+      "  File \"/home/vscode/.local/lib/python3.9/site-packages/pandas/core/indexing.py\", line 1500, in _validate_integer\n",
       "    raise IndexError(\"single positional indexer is out-of-bounds\")\n",
       "IndexError: single positional indexer is out-of-bounds\n",
       "\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
       "I apologize for the error. It seems that there is an issue with the data for the \"FB\" symbol. Let's add a try-except block to handle this issue. Please execute the following Python code:\n",
       "\n",
@@ -439,83 +443,56 @@
       "        print(f\"{symbol}: Data not available\")\n",
       "```\n",
       "\n",
-      "This code will handle any issues with the data and display \"Data not available\" for any symbols that encounter an error.\n",
+      "This code adds error handling to the `get_ytd_gain` function, so if there is an issue with a specific symbol, it will print an error message and continue with the remaining symbols.\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> USING AUTO REPLY FOR THE USER...\n",
-      "user_proxy (to assistant):\n",
+      "\u001b[31m\n",
+      ">>>>>>>> USING AUTO REPLY...\u001b[0m\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
       "exitcode: 0 (execution succeeded)\n",
       "Code output: \n",
       "Error fetching data for FB: single positional indexer is out-of-bounds\n",
-      "AAPL: 55.26%\n",
-      "MSFT: 47.19%\n",
-      "GOOGL: 37.13%\n",
-      "GOOG: 36.89%\n",
-      "AMZN: 50.47%\n",
+      "AAPL: 57.03%\n",
+      "MSFT: 41.90%\n",
+      "GOOGL: 48.77%\n",
+      "GOOG: 48.28%\n",
+      "AMZN: 54.06%\n",
       "FB: Data not available\n",
-      "TSLA: 145.40%\n",
-      "NVDA: 219.19%\n",
-      "ASML: 25.70%\n",
-      "ADBE: 55.92%\n",
-      "\n",
-      "\n",
-      "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
-      "\n",
-      "I'm glad the code executed successfully. Here are the YTD gains for the 10 largest technology companies as of today:\n",
-      "\n",
-      "- AAPL: 55.26%\n",
-      "- MSFT: 47.19%\n",
-      "- GOOGL: 37.13%\n",
-      "- GOOG: 36.89%\n",
-      "- AMZN: 50.47%\n",
-      "- FB: Data not available\n",
-      "- TSLA: 145.40%\n",
-      "- NVDA: 219.19%\n",
-      "- ASML: 25.70%\n",
-      "- ADBE: 55.92%\n",
-      "\n",
-      "Please note that the data for \"FB\" is not available due to an issue with the data source. You may want to check an alternative data source or try again later for the \"FB\" symbol.\n",
-      "\n",
-      "If you have any other questions or need further assistance, feel free to ask. Otherwise, if everything is done, please type \"TERMINATE\".\n",
+      "TSLA: 146.48%\n",
+      "NVDA: 226.67%\n",
+      "ASML: 31.38%\n",
+      "ADBE: 56.97%\n",
       "\n",
-      "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> USING AUTO REPLY FOR THE USER...\n",
-      "user_proxy (to assistant):\n",
-      "\n",
-      "\n",
-      "\n",
-      "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
-      "\n",
-      "It seems you didn't provide any input. If you have any other questions or need further assistance, feel free to ask. Otherwise, if everything is done, please type \"TERMINATE\".\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
-      ">>>>>>>> USING AUTO REPLY FOR THE USER...\n",
-      "user_proxy (to assistant):\n",
+      "I'm glad the code executed successfully. Here are the YTD gains for the 10 largest technology companies, except for Facebook (FB), for which data is not available:\n",
       "\n",
+      "- AAPL: 57.03%\n",
+      "- MSFT: 41.90%\n",
+      "- GOOGL: 48.77%\n",
+      "- GOOG: 48.28%\n",
+      "- AMZN: 54.06%\n",
+      "- TSLA: 146.48%\n",
+      "- NVDA: 226.67%\n",
+      "- ASML: 31.38%\n",
+      "- ADBE: 56.97%\n",
       "\n",
-      "\n",
-      "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
-      "\n",
-      "I understand that you might be busy or have no further questions. If you need any assistance in the future, feel free to ask. Have a great day!\n",
+      "Please note that these values are subject to change as the market fluctuates. If you need updated values, you can re-run the code.\n",
       "\n",
       "TERMINATE\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "user_proxy (to assistant):\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
-      "FB -> META\n",
+      "FB->META\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
-      "I apologize for the confusion earlier. The \"FB\" symbol has been changed to \"META\" due to the company's rebranding. Let's update the code with the new symbol and execute it again. Please run the following Python code:\n",
+      "I apologize for the confusion. Facebook has changed its stock symbol from \"FB\" to \"META\" due to the company's rebranding. Let's update the code with the new symbol and execute it again:\n",
       "\n",
       "```python\n",
       "import yfinance as yf\n",
@@ -546,72 +523,50 @@
       "        print(f\"{symbol}: Data not available\")\n",
       "```\n",
       "\n",
-      "This code will now use the \"META\" symbol instead of \"FB\" and should provide the YTD gain for Meta Platforms Inc. (formerly Facebook).\n",
+      "This code replaces \"FB\" with \"META\" in the list of top technology companies and calculates the YTD gains for the updated list.\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> USING AUTO REPLY FOR THE USER...\n",
-      "user_proxy (to assistant):\n",
+      "\u001b[31m\n",
+      ">>>>>>>> USING AUTO REPLY...\u001b[0m\n",
+      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
       "\n",
       "exitcode: 0 (execution succeeded)\n",
       "Code output: \n",
-      "AAPL: 55.26%\n",
-      "MSFT: 47.19%\n",
-      "GOOGL: 37.13%\n",
-      "GOOG: 36.89%\n",
-      "AMZN: 50.47%\n",
-      "META: 136.07%\n",
-      "TSLA: 145.40%\n",
-      "NVDA: 219.19%\n",
-      "ASML: 25.70%\n",
-      "ADBE: 55.92%\n",
-      "\n",
-      "\n",
-      "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
-      "\n",
-      "Great! The code executed successfully, and here are the updated YTD gains for the 10 largest technology companies as of today:\n",
-      "\n",
-      "- AAPL: 55.26%\n",
-      "- MSFT: 47.19%\n",
-      "- GOOGL: 37.13%\n",
-      "- GOOG: 36.89%\n",
-      "- AMZN: 50.47%\n",
-      "- META: 136.07% (formerly FB)\n",
-      "- TSLA: 145.40%\n",
-      "- NVDA: 219.19%\n",
-      "- ASML: 25.70%\n",
-      "- ADBE: 55.92%\n",
-      "\n",
-      "If you have any other questions or need further assistance, feel free to ask. Otherwise, if everything is done, please type \"TERMINATE\".\n",
-      "\n",
-      "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> USING AUTO REPLY FOR THE USER...\n",
-      "user_proxy (to assistant):\n",
-      "\n",
+      "AAPL: 57.03%\n",
+      "MSFT: 41.90%\n",
+      "GOOGL: 48.77%\n",
+      "GOOG: 48.28%\n",
+      "AMZN: 54.06%\n",
+      "META: 160.93%\n",
+      "TSLA: 146.48%\n",
+      "NVDA: 226.67%\n",
+      "ASML: 31.38%\n",
+      "ADBE: 56.97%\n",
       "\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
+      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
       "\n",
-      "It seems you didn't provide any input. If you have any other questions or need further assistance, feel free to ask. Otherwise, if everything is done, please type \"TERMINATE\".\n",
+      "Great! The code executed successfully, and here are the updated YTD gains for the 10 largest technology companies, including Meta Platforms (formerly Facebook):\n",
       "\n",
-      "--------------------------------------------------------------------------------\n",
-      "\n",
-      ">>>>>>>> USING AUTO REPLY FOR THE USER...\n",
-      "user_proxy (to assistant):\n",
+      "- AAPL: 57.03%\n",
+      "- MSFT: 41.90%\n",
+      "- GOOGL: 48.77%\n",
+      "- GOOG: 48.28%\n",
+      "- AMZN: 54.06%\n",
+      "- META: 160.93%\n",
+      "- TSLA: 146.48%\n",
+      "- NVDA: 226.67%\n",
+      "- ASML: 31.38%\n",
+      "- ADBE: 56.97%\n",
       "\n",
-      "\n",
-      "\n",
-      "--------------------------------------------------------------------------------\n",
-      "assistant (to user_proxy):\n",
-      "\n",
-      "I understand that you might be busy or have no further questions. If you need any assistance in the future, feel free to ask. Have a great day!\n",
+      "Please note that these values are subject to change as the market fluctuates. If you need updated values, you can re-run the code.\n",
       "\n",
       "TERMINATE\n",
       "\n",
-      "--------------------------------------------------------------------------------\n"
+      "--------------------------------------------------------------------------------\n",
+      "\u001b[31m\n",
+      ">>>>>>>> NO HUMAN INPUT RECEIVED.\u001b[0m\n"
      ]
     }
    ],
diff --git a/notebook/autogen_chatgpt_gpt4.ipynb b/notebook/autogen_chatgpt_gpt4.ipynb
index 1417d795db..463c559301 100644
--- a/notebook/autogen_chatgpt_gpt4.ipynb
+++ b/notebook/autogen_chatgpt_gpt4.ipynb
@@ -31,7 +31,7 @@
     "\n",
     "## Requirements\n",
     "\n",
-    "FLAML requires `Python>=3.7`. To run this notebook example, please install flaml with the [openai,blendsearch] option:\n",
+    "FLAML requires `Python>=3.8`. To run this notebook example, please install flaml with the [openai,blendsearch] option:\n",
     "```bash\n",
     "pip install flaml[openai,blendsearch]\n",
     "```"
@@ -58,9 +58,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Import the oai and tune subpackages from flaml\n",
-    "\n",
-    "FLAML has provided an API for hyperparameter optimization of OpenAI ChatGPT models: `oai.ChatCompletion.tune` and to make a request with the tuned config: `oai.ChatCompletion.create`. First, we import oai from flaml:"
+    "FLAML has provided an API for hyperparameter optimization of OpenAI ChatGPT models: `autogen.ChatCompletion.tune` and to make a request with the tuned config: `autogen.ChatCompletion.create`. First, we import autogen from flaml:"
    ]
   },
   {
@@ -79,7 +77,7 @@
    },
    "outputs": [],
    "source": [
-    "from flaml import oai"
+    "from flaml import autogen"
    ]
   },
   {
@@ -111,7 +109,7 @@
    },
    "outputs": [],
    "source": [
-    "config_list = oai.config_list_openai_aoai()"
+    "config_list = autogen.config_list_openai_aoai()"
    ]
   },
   {
@@ -360,7 +358,7 @@
    },
    "outputs": [],
    "source": [
-    "oai.ChatCompletion.set_cache(seed)"
+    "autogen.ChatCompletion.set_cache(seed)"
    ]
   },
   {
@@ -443,7 +441,7 @@
     "import logging\n",
     "\n",
     "prompts = [\"{problem} Solve the problem carefully. Simplify your answer as much as possible. Put the final answer in \\\\boxed{{}}.\"]\n",
-    "config, analysis = oai.ChatCompletion.tune(\n",
+    "config, analysis = autogen.ChatCompletion.tune(\n",
     "    data=tune_data,  # the data for tuning\n",
     "    metric=\"success_vote\",  # the metric to optimize\n",
     "    mode=\"max\",  # the optimization mode\n",
@@ -952,8 +950,8 @@
     }
    ],
    "source": [
-    "response = oai.ChatCompletion.create(context=tune_data[1], config_list=config_list, **config)\n",
-    "metric_results = eval_math_responses(oai.ChatCompletion.extract_text(response), **tune_data[1])\n",
+    "response = autogen.ChatCompletion.create(context=tune_data[1], config_list=config_list, **config)\n",
+    "metric_results = eval_math_responses(autogen.ChatCompletion.extract_text(response), **tune_data[1])\n",
     "print(\"response on an example data instance:\", response)\n",
     "print(\"metric_results on the example data instance:\", metric_results)\n"
    ]
@@ -965,7 +963,7 @@
    "source": [
     "### Evaluate the success rate on the test data\n",
     "\n",
-    "You can use flaml's `oai.ChatCompletion.test` to evaluate the performance of an entire dataset with the tuned config. The following code will take a while (30 mins to 1 hour) to evaluate all the test data instances if uncommented and run. It will cost roughly $3. "
+    "You can use `autogen.ChatCompletion.test` to evaluate the performance of an entire dataset with the tuned config. The following code will take a while (30 mins to 1 hour) to evaluate all the test data instances if uncommented and run. It will cost roughly $3. "
    ]
   },
   {
@@ -1190,7 +1188,7 @@
     }
    ],
    "source": [
-    "# result = oai.ChatCompletion.test(test_data, logging_level=logging.INFO, config_list=config_list, **config)\n",
+    "# result = autogen.ChatCompletion.test(test_data, logging_level=logging.INFO, config_list=config_list, **config)\n",
     "# print(\"performance on test data with the tuned config:\", result)"
    ]
   },
@@ -1219,7 +1217,7 @@
     "# the following code will cost roughly $2 if uncommented and run.\n",
     "\n",
     "# default_config = {\"model\": 'gpt-4', \"prompt\": prompts[0]}\n",
-    "# default_result = oai.ChatCompletion.test(test_data, config_list=config_list, **default_config)\n",
+    "# default_result = autogen.ChatCompletion.test(test_data, config_list=config_list, **default_config)\n",
     "# print(\"performance on test data from gpt-4 with a default config:\", default_result)"
    ]
   },
@@ -1259,7 +1257,7 @@
     "# The following evaluation costs $3 and longer than one hour if you uncomment it and run it.\n",
     "\n",
     "# config_n2 = {\"model\": 'gpt-4', \"prompt\": prompts[0], \"n\": 2}\n",
-    "# result_n2 = oai.ChatCompletion.test(test_data, config_list=config_list, **config_n2)\n",
+    "# result_n2 = autogen.ChatCompletion.test(test_data, config_list=config_list, **config_n2)\n",
     "# print(\"performance on test data from gpt-4 with a default config and n=2:\", result_n2)\n"
    ]
   },
@@ -1280,7 +1278,7 @@
     "# The following evaluation costs $8 and longer than one hour if you uncomment it and run it.\n",
     "\n",
     "# config_n5 = {\"model\": 'gpt-4', \"prompt\": prompts[0], \"n\": 5}\n",
-    "# result_n5 = oai.ChatCompletion.test(test_data, config_list=config_list, **config_n5)\n",
+    "# result_n5 = autogen.ChatCompletion.test(test_data, config_list=config_list, **config_n5)\n",
     "# print(\"performance on test data from gpt-4 with a default config and n=5:\", result_n5)"
    ]
   },
diff --git a/notebook/autogen_openai_completion.ipynb b/notebook/autogen_openai_completion.ipynb
index a418db4e8f..225bb1fcc8 100644
--- a/notebook/autogen_openai_completion.ipynb
+++ b/notebook/autogen_openai_completion.ipynb
@@ -29,7 +29,7 @@
     "\n",
     "## Requirements\n",
     "\n",
-    "FLAML requires `Python>=3.7`. To run this notebook example, please install flaml with the [autogen,blendsearch] option:\n",
+    "FLAML requires `Python>=3.8`. To run this notebook example, please install flaml with the [autogen,blendsearch] option:\n",
     "```bash\n",
     "pip install flaml[autogen,blendsearch]\n",
     "```"
@@ -37,7 +37,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 1,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2023-02-24T23:25:36.910966Z",
@@ -48,7 +48,7 @@
    },
    "outputs": [],
    "source": [
-    "# %pip install flaml[autogen,blendsearch] datasets"
+    "# %pip install flaml[autogen,blendsearch]~=2.0.0rc4 datasets"
    ]
   },
   {
@@ -69,7 +69,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 2,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2023-02-24T23:25:36.917301Z",
@@ -80,9 +80,9 @@
    },
    "outputs": [],
    "source": [
-    "from flaml import oai\n",
+    "from flaml import autogen\n",
     "\n",
-    "endpoint_list = oai.config_list_openai_aoai()\n",
+    "endpoint_list = autogen.config_list_openai_aoai()\n",
     "# the endpoint_list looks like this:\n",
     "# endpoint_list = [\n",
     "#     {\n",
@@ -102,7 +102,7 @@
     "#     },  # another Azure OpenAI API endpoint for gpt-4\n",
     "# ]\n",
     "\n",
-    "config_list = oai.config_list_from_json(\n",
+    "config_list = autogen.config_list_from_json(\n",
     "    env_or_file=\"OAI_CONFIG_LIST\",\n",
     "    filter_dict={\n",
     "        \"model\": {\n",
@@ -151,7 +151,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 3,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2023-02-24T23:25:36.931255Z",
@@ -171,7 +171,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "95bea2ebf248403aac1ee64e8bebd5e3",
+       "model_id": "8e08cc907707418a86a3da668e45326b",
        "version_major": 2,
        "version_minor": 0
       },
@@ -228,7 +228,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 4,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2023-02-24T23:25:39.152156Z",
@@ -280,7 +280,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 5,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2023-02-24T23:25:39.158398Z",
@@ -325,7 +325,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 6,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2023-02-24T23:25:39.164187Z",
@@ -337,11 +337,10 @@
    "outputs": [],
    "source": [
     "from functools import partial\n",
-    "from flaml.autogen.code_utils import eval_function_completions, generate_assertions\n",
     "\n",
     "eval_with_generated_assertions = partial(\n",
-    "    eval_function_completions,\n",
-    "    assertions=partial(generate_assertions, config_list=config_list),\n",
+    "    autogen.code_utils.eval_function_completions,\n",
+    "    assertions=partial(autogen.code_utils.generate_assertions, config_list=config_list),\n",
     "    use_docker=False,\n",
     "    # Please set use_docker=True if you have docker available to run the generated code.\n",
     "    # Using docker is safer than running the generated code directly.\n",
@@ -361,16 +360,14 @@
     "\n",
     "## Use the tuning data to find a good configuration\n",
     "\n",
-    "### Import the oai and tune subpackages from flaml.\n",
-    "\n",
-    "FLAML has provided an API for hyperparameter optimization of OpenAI models: `oai.Completion.tune` and to make a request with the tuned config: `oai.Completion.create`.\n",
+    "FLAML has provided an API for hyperparameter optimization of OpenAI models: `autogen.Completion.tune` and to make a request with the tuned config: `autogen.Completion.create`.\n",
     "\n",
     "For (local) reproducibility and cost efficiency, we cache responses from OpenAI with a controllable seed."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 7,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2023-02-24T23:25:40.587815Z",
@@ -384,7 +381,7 @@
    },
    "outputs": [],
    "source": [
-    "oai.Completion.set_cache(seed)"
+    "autogen.Completion.set_cache(seed)"
    ]
   },
   {
@@ -431,7 +428,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 8,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2023-02-24T23:25:40.593603Z",
@@ -445,30 +442,30 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[32m[I 2023-07-18 17:50:33,625]\u001b[0m A new study created in memory with name: optuna\u001b[0m\n",
-      "\u001b[32m[I 2023-07-18 17:50:33,628]\u001b[0m A new study created in memory with name: optuna\u001b[0m\n"
+      "\u001b[32m[I 2023-07-30 04:19:08,150]\u001b[0m A new study created in memory with name: optuna\u001b[0m\n",
+      "\u001b[32m[I 2023-07-30 04:19:08,153]\u001b[0m A new study created in memory with name: optuna\u001b[0m\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "[flaml.tune.tune: 07-18 17:50:34] {805} INFO - trial 1 config: {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-ada-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}\n",
-      "[flaml.tune.tune: 07-18 17:51:28] {197} INFO - result: {'index_selected': 26.0, 'succeed_assertions': 0.0, 'success': 0.0, 'gen_cost': 0.0004240249999999999, 'assertions': 'assert vowels_count(\"abcde\") == 2\\nassert vowels_count(\"ACEDY\") == 3', 'total_cost': 0.0104828, 'cost': 0.0104828, 'inference_cost': 0.00023374, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-ada-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'text-ada-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 54.489821434020996}\n",
-      "[flaml.tune.tune: 07-18 17:51:28] {805} INFO - trial 2 config: {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-babbage-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}\n",
-      "[flaml.tune.tune: 07-18 17:52:11] {197} INFO - result: {'index_selected': 26.0, 'succeed_assertions': 0.0, 'success': 0.0, 'gen_cost': 0.0004240249999999999, 'assertions': 'assert vowels_count(\"abcde\") == 2\\nassert vowels_count(\"ACEDY\") == 3', 'total_cost': 0.030993300000000005, 'cost': 0.020510499999999997, 'inference_cost': 0.001025525, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-babbage-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'text-babbage-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 42.71049189567566}\n",
-      "[flaml.tune.tune: 07-18 17:52:11] {805} INFO - trial 3 config: {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-davinci-003', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}\n",
-      "[flaml.tune.tune: 07-18 17:54:24] {197} INFO - result: {'index_selected': 3.65, 'succeed_assertions': 0.9, 'success': 0.7, 'gen_cost': 0.0004240249999999999, 'assertions': 'assert vowels_count(\"abcde\") == 2\\nassert vowels_count(\"ACEDY\") == 3', 'total_cost': 0.8669533, 'cost': 0.83596, 'inference_cost': 0.041073000000000005, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-davinci-003', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'text-davinci-003', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 133.33429384231567}\n",
-      "[flaml.tune.tune: 07-18 17:54:24] {805} INFO - trial 4 config: {'prompt': 1, 'stop': 0, 'subspace': {'model': 'gpt-3.5-turbo', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}\n",
-      "[flaml.tune.tune: 07-18 17:55:48] {197} INFO - result: {'index_selected': 16.65, 'succeed_assertions': 0.4, 'success': 0.35, 'gen_cost': 0.0004240249999999999, 'assertions': 'assert vowels_count(\"abcde\") == 2\\nassert vowels_count(\"ACEDY\") == 3', 'total_cost': 0.9311253000000002, 'cost': 0.064172, 'inference_cost': 0.0031784499999999998, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'gpt-3.5-turbo', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'gpt-3.5-turbo', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 83.27373838424683}\n",
-      "[flaml.tune.tune: 07-18 17:55:48] {805} INFO - trial 5 config: {'prompt': 1, 'stop': 0, 'subspace': {'model': 'gpt-4', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}\n",
-      "[flaml.tune.tune: 07-18 17:56:06] {197} INFO - result: {'success': 0, 'total_cost': 1.0261653000000002, 'cost': 0.09504, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'gpt-4', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'gpt-4', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 18.89783787727356}\n",
-      "[flaml.tune.tune: 07-18 17:56:06] {828} WARNING - fail to sample a trial for 100 times in a row, stopping.\n"
+      "[flaml.tune.tune: 07-30 04:19:08] {805} INFO - trial 1 config: {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-ada-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}\n",
+      "[flaml.tune.tune: 07-30 04:22:35] {197} INFO - result: {'index_selected': 26.0, 'succeed_assertions': 0.0, 'success': 0.0, 'gen_cost': 0.000460625, 'assertions': 'assert vowels_count(\"abcde\") == 2\\nassert vowels_count(\"ACEDY\") == 3', 'total_cost': 0.010514800000000003, 'cost': 0.010514800000000003, 'inference_cost': 0.00023534000000000003, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-ada-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'text-ada-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 207.29033374786377}\n",
+      "[flaml.tune.tune: 07-30 04:22:35] {805} INFO - trial 2 config: {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-babbage-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}\n",
+      "[flaml.tune.tune: 07-30 04:23:18] {197} INFO - result: {'index_selected': 26.0, 'succeed_assertions': 0.0, 'success': 0.0, 'gen_cost': 0.000460625, 'assertions': 'assert vowels_count(\"abcde\") == 2\\nassert vowels_count(\"ACEDY\") == 3', 'total_cost': 0.0300243, 'cost': 0.019509500000000003, 'inference_cost': 0.0009754750000000001, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-babbage-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'text-babbage-001', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 42.417603969573975}\n",
+      "[flaml.tune.tune: 07-30 04:23:18] {805} INFO - trial 3 config: {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-davinci-003', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}\n",
+      "[flaml.tune.tune: 07-30 04:24:20] {197} INFO - result: {'index_selected': 2.35, 'succeed_assertions': 0.95, 'success': 0.65, 'gen_cost': 0.000460625, 'assertions': 'assert vowels_count(\"abcde\") == 2\\nassert vowels_count(\"ACEDY\") == 3', 'total_cost': 0.8658043000000002, 'cost': 0.8357800000000002, 'inference_cost': 0.04093000000000001, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-davinci-003', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'text-davinci-003', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 62.81497287750244}\n",
+      "[flaml.tune.tune: 07-30 04:24:20] {805} INFO - trial 4 config: {'prompt': 1, 'stop': 0, 'subspace': {'model': 'gpt-3.5-turbo', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}\n",
+      "[flaml.tune.tune: 07-30 04:25:39] {197} INFO - result: {'index_selected': 13.95, 'succeed_assertions': 0.55, 'success': 0.5, 'gen_cost': 0.000460625, 'assertions': 'assert vowels_count(\"abcde\") == 2\\nassert vowels_count(\"ACEDY\") == 3', 'total_cost': 0.9462703000000001, 'cost': 0.08046600000000001, 'inference_cost': 0.00399515, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'gpt-3.5-turbo', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'gpt-3.5-turbo', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 79.03474521636963}\n",
+      "[flaml.tune.tune: 07-30 04:25:39] {805} INFO - trial 5 config: {'prompt': 1, 'stop': 0, 'subspace': {'model': 'gpt-4', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}\n",
+      "[flaml.tune.tune: 07-30 04:25:50] {197} INFO - result: {'success': 0, 'total_cost': 1.0053703, 'cost': 0.0591, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'gpt-4', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'gpt-4', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 10.245523691177368}\n",
+      "[flaml.tune.tune: 07-30 04:25:50] {828} WARNING - fail to sample a trial for 100 times in a row, stopping.\n"
      ]
     }
    ],
    "source": [
-    "config, analysis = oai.Completion.tune(\n",
+    "config, analysis = autogen.Completion.tune(\n",
     "    data=tune_data,  # the data for tuning\n",
     "    metric=\"success\",  # the metric to optimize\n",
     "    mode=\"max\",  # the optimization mode\n",
@@ -496,12 +493,12 @@
    "source": [
     "### Output tuning results\n",
     "\n",
-    "After the tuning, we can print out the config and the result found by FLAML:"
+    "After the tuning, we can print out the config and the result found by autogen:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 9,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2023-02-24T23:26:38.352710Z",
@@ -516,7 +513,7 @@
      "output_type": "stream",
      "text": [
       "optimized config {'prompt': '# Python 3{definition}', 'stop': ['\\nclass', '\\ndef', '\\nif', '\\nprint'], 'model': 'text-davinci-003', 'max_tokens': 148, 'n': 27, 'top_p': 0.755486898036596}\n",
-      "best result on tuning data {'index_selected': 3.65, 'succeed_assertions': 0.9, 'success': 0.7, 'gen_cost': 0.0004240249999999999, 'assertions': 'assert vowels_count(\"abcde\") == 2\\nassert vowels_count(\"ACEDY\") == 3', 'total_cost': 0.8669533, 'cost': 0.83596, 'inference_cost': 0.041073000000000005, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-davinci-003', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'text-davinci-003', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 133.33429384231567}\n"
+      "best result on tuning data {'index_selected': 2.35, 'succeed_assertions': 0.95, 'success': 0.65, 'gen_cost': 0.000460625, 'assertions': 'assert vowels_count(\"abcde\") == 2\\nassert vowels_count(\"ACEDY\") == 3', 'total_cost': 0.8658043000000002, 'cost': 0.8357800000000002, 'inference_cost': 0.04093000000000001, 'training_iteration': 0, 'config': {'prompt': 1, 'stop': 0, 'subspace': {'model': 'text-davinci-003', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}}, 'config/prompt': 1, 'config/stop': 0, 'config/subspace': {'model': 'text-davinci-003', 'max_tokens': 148, 'temperature_or_top_p': {'top_p': 0.755486898036596}, 'n': 27}, 'experiment_tag': 'exp', 'time_total_s': 62.81497287750244}\n"
      ]
     }
    ],
@@ -541,7 +538,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 10,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2023-02-24T23:26:38.359902Z",
@@ -560,37 +557,38 @@
      "output_type": "stream",
      "text": [
       "{\n",
-      "  \"id\": \"cmpl-7djA2WbYCudgjWodOG8aAHyEZ2lex\",\n",
+      "  \"id\": \"cmpl-7hsFhPX6faeWYaT4y3C7IkQAgNbZR\",\n",
+      "  \"warning\": \"This model version is deprecated. Migrate before January 4, 2024 to avoid disruption of service. Learn more https://platform.openai.com/docs/deprecations\",\n",
       "  \"object\": \"text_completion\",\n",
-      "  \"created\": 1689702746,\n",
+      "  \"created\": 1690691005,\n",
       "  \"model\": \"text-davinci-003\",\n",
       "  \"choices\": [\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i]-guess[i]))\\n    return result\",\n",
+      "      \"text\": \"    results = []\\n    for i in range(len(game)):\\n        if game[i] == guess[i]:\\n            results.append(0)\\n        else:\\n            results.append(abs(game[i]-guess[i]))\\n    return results\",\n",
       "      \"index\": 0,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        if game[i] == guess[i]:\\n            result.append(0)\\n        else:\\n            result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
       "      \"index\": 1,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        if game[i] == guess[i]:\\n            result.append(0)\\n        else:\\n            result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i]-guess[i]))\\n    return result\",\n",
       "      \"index\": 2,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
+      "      \"text\": \"    res = []\\n    for i in range(len(game)):\\n        res.append(abs(game[i]-guess[i]))\\n    return res\",\n",
       "      \"index\": 3,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    results = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        results.append(diff)\\n    return results\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
       "      \"index\": 4,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
@@ -602,115 +600,115 @@
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    # write your code here\\n    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
       "      \"index\": 6,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        if game[i] == guess[i]:\\n            result.append(0)\\n        else:\\n            result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
+      "      \"text\": \"    results = []\\n    for i in range(len(game)):\\n        results.append(abs(game[i] - guess[i]))\\n    return results\",\n",
       "      \"index\": 7,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
+      "      \"text\": \"    res = []\\n    for i in range(len(game)):\\n        res.append(abs(game[i]-guess[i]))\\n    return res\",\n",
       "      \"index\": 8,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i]-guess[i]))\\n    return result\",\n",
       "      \"index\": 9,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        if game[i] == guess[i]:\\n            result.append(0)\\n        else:\\n            result.append(abs(game[i]-guess[i]))\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
       "      \"index\": 10,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
       "      \"index\": 11,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        if game[i] == guess[i]:\\n            result.append(0)\\n        else:\\n            result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
       "      \"index\": 12,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
+      "      \"text\": \"    # set up empty list to store differences\\n    diff = []\\n    # iterate through the game list and guess list\\n    for i in range(len(game)):\\n        # check if the guess is equal to the game\\n        if game[i] == guess[i]:\\n            # if so, append 0 to the diff list\\n            diff.append(0)\\n        # otherwise, calculate the difference between the guess and the game\\n        else:\\n            diff.append(abs(game[i]-guess[i]))\\n    # return the diff list\\n    return diff\",\n",
       "      \"index\": 13,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    #initialize empty list\\n    result = []\\n    #iterate over the two lists\\n    for i in range(len(game)):\\n        #append the absolute difference between the guess and the score to the result list\\n        result.append(abs(game[i]-guess[i]))\\n    #return the result list\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i]-guess[i]))\\n    return result\",\n",
       "      \"index\": 14,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        if game[i] == guess[i]:\\n            result.append(0)\\n        else:\\n            result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
       "      \"index\": 15,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
       "      \"index\": 16,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
       "      \"index\": 17,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
       "      \"index\": 18,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
       "      \"index\": 19,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
       "      \"index\": 20,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    res = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        res.append(diff)\\n    return res\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
       "      \"index\": 21,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
       "      \"index\": 22,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
+      "      \"text\": \"    # your code here\\n    result = []\\n    for i in range(len(game)):\\n        if game[i] == guess[i]:\\n            result.append(0)\\n        else:\\n            result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
       "      \"index\": 23,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        result.append(diff)\\n    return result\",\n",
+      "      \"text\": \"    results = []\\n    for i in range(len(game)):\\n        diff = abs(game[i] - guess[i])\\n        results.append(diff)\\n    return results\",\n",
       "      \"index\": 24,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
@@ -722,7 +720,7 @@
       "      \"finish_reason\": \"stop\"\n",
       "    },\n",
       "    {\n",
-      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i]-guess[i]))\\n    return result\",\n",
+      "      \"text\": \"    result = []\\n    for i in range(len(game)):\\n        result.append(abs(game[i] - guess[i]))\\n    return result\",\n",
       "      \"index\": 26,\n",
       "      \"logprobs\": null,\n",
       "      \"finish_reason\": \"stop\"\n",
@@ -730,21 +728,21 @@
       "  ],\n",
       "  \"usage\": {\n",
       "    \"prompt_tokens\": 243,\n",
-      "    \"completion_tokens\": 1244,\n",
-      "    \"total_tokens\": 1487\n",
+      "    \"completion_tokens\": 1264,\n",
+      "    \"total_tokens\": 1507\n",
       "  },\n",
-      "  \"cost\": 0.029740000000000003,\n",
+      "  \"cost\": 0.03014,\n",
       "  \"config_id\": 0,\n",
       "  \"pass_filter\": true\n",
       "}\n",
-      "{'index_selected': 0, 'succeed_assertions': True, 'success': True, 'gen_cost': 0.0005660000000000001, 'assertions': 'assert compare([1,2,3,4,5,1],[1,2,3,4,2,-2]) == [0,0,0,0,3,3]\\nassert compare([0,5,0,0,0,4],[4,1,1,0,0,-2]) == [4,4,1,0,0,6]'}\n"
+      "{'index_selected': 0, 'succeed_assertions': True, 'success': True, 'gen_cost': 0.000702, 'assertions': 'assert compare([1,2,3,4,5,1],[1,2,3,4,2,-2]) == [0,0,0,0,3,3]\\nassert compare([0,5,0,0,0,4],[4,1,1,0,0,-2]) == [4,4,1,0,0,6]'}\n"
      ]
     }
    ],
    "source": [
-    "response = oai.Completion.create(context=tune_data[1], config_list=endpoint_list, **config)\n",
+    "response = autogen.Completion.create(context=tune_data[1], config_list=endpoint_list, **config)\n",
     "print(response)\n",
-    "print(eval_with_generated_assertions(oai.Completion.extract_text(response), **tune_data[1]))\n"
+    "print(eval_with_generated_assertions(autogen.Completion.extract_text(response), **tune_data[1]))\n"
    ]
   },
   {
@@ -754,12 +752,12 @@
    "source": [
     "### Evaluate the success rate on the test data\n",
     "\n",
-    "You can use flaml's `oai.Completion.test` to evaluate the performance of an entire dataset with the tuned config. The following code will take a while to evaluate all the 144 test data instances. The cost is about $6 if you uncomment it and run it."
+    "You can use `autogen.Completion.test` to evaluate the performance of an entire dataset with the tuned config. The following code will take a while to evaluate all the 144 test data instances. The cost is about $6 if you uncomment it and run it."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 12,
    "metadata": {
     "execution": {
      "iopub.execute_input": "2023-02-24T23:26:39.347295Z",
@@ -773,12 +771,12 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "performance on test data with the tuned config: {'index_selected': 5.208333333333333, 'succeed_assertions': 0.8402777777777778, 'success': 0.7777777777777778, 'gen_cost': 0.00045375000000000005, 'cost': 5.785519999999999, 'inference_cost': 0.04017722222222222}\n"
+      "performance on test data with the tuned config: {'index_selected': 5.222222222222222, 'succeed_assertions': 0.8402777777777778, 'success': 0.7569444444444444, 'gen_cost': 0.00044632638888888885, 'cost': 5.704979999999999, 'inference_cost': 0.03961791666666666}\n"
      ]
     }
    ],
    "source": [
-    "# result = oai.Completion.test(test_data, **config)\n",
+    "# result = autogen.Completion.test(test_data, config_list=endpoint_list, **config)\n",
     "# print(\"performance on test data with the tuned config:\", result)"
    ]
   },
diff --git a/notebook/research/autogen_code.ipynb b/notebook/research/autogen_code.ipynb
index 4fb0dfd8b0..653bd94393 100644
--- a/notebook/research/autogen_code.ipynb
+++ b/notebook/research/autogen_code.ipynb
@@ -170,7 +170,7 @@
    "outputs": [],
    "source": [
     "from flaml.autogen.code_utils import eval_function_completions, implement\n",
-    "from flaml import oai"
+    "from flaml.autogen import oai"
    ]
   },
   {
@@ -409,7 +409,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.16 (main, Apr 12 2023, 14:54:44) \n[GCC 10.2.1 20210110]"
+   "version": "3.9.16"
   },
   "vscode": {
    "interpreter": {
diff --git a/notebook/research/math_level5counting.ipynb b/notebook/research/math_level5counting.ipynb
index 1f27036d14..17721cbd60 100644
--- a/notebook/research/math_level5counting.ipynb
+++ b/notebook/research/math_level5counting.ipynb
@@ -233,7 +233,7 @@
     }
    },
    "source": [
-    "### Import the oai and tune subpackages from flaml.\n"
+    "### Import the oai subpackage from flaml.\n"
    ]
   },
   {
@@ -252,7 +252,7 @@
    },
    "outputs": [],
    "source": [
-    "from flaml import oai"
+    "from flaml.autogen import oai"
    ]
   },
   {
@@ -313,7 +313,7 @@
    "source": [
     "### Evaluate the success rate on the test data\n",
     "\n",
-    "You can use flaml's `oai.ChatCompletion.test` to evaluate the performance of an entire dataset with a config."
+    "You can use `oai.ChatCompletion.test` to evaluate the performance of an entire dataset with a config."
    ]
   },
   {
diff --git a/setup.py b/setup.py
index b7ddbc0c23..65696cda19 100644
--- a/setup.py
+++ b/setup.py
@@ -139,8 +139,8 @@
         ],
         "benchmark": ["catboost>=0.26", "psutil==5.8.0", "xgboost==1.3.3", "pandas==1.1.4"],
         "openai": ["openai==0.27.8", "diskcache"],
-        "autogen": ["openai==0.27.8", "diskcache", "docker"],
-        "mathchat": ["openai==0.27.8", "diskcache", "docker", "sympy", "pydantic==1.10.9", "wolframalpha"],
+        "autogen": ["openai==0.27.8", "diskcache", "termcolor"],
+        "mathchat": ["openai==0.27.8", "diskcache", "termcolor", "sympy", "pydantic==1.10.9", "wolframalpha"],
         "synapse": [
             "joblibspark>=0.5.0",
             "optuna==2.8.0",
diff --git a/website/blog/2023-06-28-MathChat/img/mathchatflow.png b/website/blog/2023-06-28-MathChat/img/mathchatflow.png
new file mode 100644
index 0000000000..19b41ead78
Binary files /dev/null and b/website/blog/2023-06-28-MathChat/img/mathchatflow.png differ
diff --git a/website/blog/2023-06-28-MathChat/img/result.png b/website/blog/2023-06-28-MathChat/img/result.png
new file mode 100644
index 0000000000..88a269e316
Binary files /dev/null and b/website/blog/2023-06-28-MathChat/img/result.png differ
diff --git a/website/blog/2023-06-28-MathChat/index.mdx b/website/blog/2023-06-28-MathChat/index.mdx
new file mode 100644
index 0000000000..0603a38983
--- /dev/null
+++ b/website/blog/2023-06-28-MathChat/index.mdx
@@ -0,0 +1,94 @@
+---
+title: MathChat - An Conversational Framework to Solve Math Problems
+authors: yiranwu
+tags: [LLM, GPT, research]
+---
+
+![MathChat WorkFlow](img/mathchatflow.png)
+**TL;DR:**
+* **We introduce MathChat, a conversational framework leveraging Large Language Models (LLMs), specifically GPT-4, to solve advanced mathematical problems.**
+* **MathChat improves LLM's performance on challenging math problem-solving, outperforming basic prompting and other strategies by about 6%. The improvement was especially notable in the Algebra category, with a 15% increase in accuracy.**
+* **Despite the advancement, GPT-4 still struggles to solve very challenging math problems, even with effective prompting strategies. Further improvements are needed, such as the development of more specific assistant models or the integration of new tools and prompts.**
+
+Recent Large Language Models (LLMs) like GTP-3.5 and GPT-4 have demonstrated astonishing abilities over previous models on various tasks, such as text generation, question answering, and code generation. Moreover, these models can communicate with humans through conversations and remember previous contexts, making it easier for humans to interact with them. These models play an increasingly important role in our daily lives assisting people with different tasks, such as writing emails, summarizing documents, and writing code.
+
+In this blog post, we probe into the problem-solving capabilities of LLMs. Specifically, we are interested in their capabilities to solve advanced math problems, which could be representative of a broader class of problems that require precise reasoning and also have deterministic solutions.
+
+We introduce MathChat, a conversational framework designed for solving challenging math problems with LLMs. This framework takes advantage of the chat-optimized feature of state-of-the-art LLMs, where a user proxy agent and an LLM assistant work together to tackle math problems. We also test previous prompting techniques for comparison.
+
+## The MathChat Framework
+
+MathChat simulates a conversation between the LLM assistant and a user proxy agent. As the name indicates, the user proxy agent acts as a proxy for the user, which is responsible for communicating with the LLM assistant and continuing the conversation in a desired manner.
+
+The proxy agent first presents a math problem to the LLM assistant, framed by an initial prompt that guides the assistant in how it should work collaboratively to solve the problem. We incorporate existing prompting methods such as CoT (chain-of-thought) and tool-using are incorporated in the initial prompt. Key components of the prompt include:
+
+- **Tool-using Prompt:** This guides the LLM assistant to use Python code in the correct format to address the problem, allowing the user proxy agent to parse the code and provide the appropriate results.
+
+- **Problem-Solving Strategy Selection Prompt:** The assistant is instructed to choose one of three potential problem-solving strategies, including:
+  1. Writing a Python program to solve the problem directly.
+  2. Solving the problem directly without using Python, which exercises LLMs' inherent reasoning capabilities.
+  3. Solving the problem step by step with Python. This should be used when the first two strategies aren't suitable.
+
+- **Final Answer Encapsulation Prompt:** This part instructs the assistant to put the final answer in *\boxed*{}.
+
+The prompt design in MathChat is flexible and can easily be refined to include other tools, such as Wolfram Alpha.
+
+Let's take a look at an example between the **User Proxy Agent** and the **LLM Assistant** (GPT-4). The conversation focuses on how to solve inequality using Python.
+(The conversation is modified for readability.)
+
+| Participant | Dialogue |
+|-------------|----------|
+| User Proxy Agent | (Initial Prompt: Let's solve a math problem with Python...) Problem: Find all x that satisfy the inequality (2x+10)(x+3) < (3x+9)(x+8). Express your answer in interval notation. |
+| LLM assistant | The key idea to solve this problem is to find the critical points, which are the points where the inequality is equal. Then, test the intervals between the critical points to find the interval for which the inequality holds. |
+| LLM assistant (continued) | Here's Python code to find the critical points: `from sympy import symbols, Eq, solve; x = symbols("x"); lhs = (2*x + 10)*(x + 3); rhs = (3*x + 9)*(x + 8); equation = Eq(lhs, rhs); critical_points = solve(equation, x); critical_points` |
+| User Proxy Agent | The critical points are: [-14, -3] |
+| LLM assistant | We found the critical points to be x = -14 and x = -3. Now we'll test the intervals between these critical points as well as points before and after the critical points. The intervals are: x < -14, -14 < x < -3, x > -3. |
+| LLM assistant (continued) | Here's Python code to test the inequality for each interval: `from sympy import symbols; x = symbols("x"); inequality = (2*x + 10)*(x + 3) < (3*x + 9)*(x + 8); test_points = [-15, -5, 0]; intervals = []; for point in test_points: if inequality.subs(x, point): intervals.append(point); intervals` |
+| User Proxy Agent | The intervals where the inequality holds true are for test points: [-15, 0] |
+| LLM assistant | We have tested the intervals and found that the inequality holds for: 1. x < -14 (test point: -15) 2. x > -3 (test point: 0). Therefore, the solution to the inequality (2x+10)(x+3) < (3x+9)(x+8) is x ∈ (-∞, -14) ∪ (-3, ∞) |
+
+## Experiment Setup
+
+We evaluate the improvement brought by MathChat.
+
+For the experiment, we focus on the level-5 problems from the MATH dataset, which are composed of high school competition problems. These problems include the application of theorems and complex equation derivation and are challenging even for undergraduate students. We evaluate 6 of 7 categories from the dataset (excluding Geometry): Prealgebra, Algebra, Number Theory, Counting and Probability, Intermediate Algebra, and Precalculus.
+
+We evaluate GPT-4 and use the default configuration of the OpenAI API. To access the final performance, we manually compare the final answer with the correct answer.  For the vanilla prompt, Program Synthesis, and MathChat, we have GPT-4 enclose the final answer in *\boxed*{}, and we take the return of the function in PoT as the final answer.
+
+
+We also evaluate the following methods for comparison:
+
+1. **Vanilla prompting:** Evaluates GPT-4's direct problem-solving capability. The prompt used is: *" Solve the problem carefully. Put the final answer in \boxed{}"*.
+
+2. **Program of Thoughts (PoT):** Uses a zero-shot PoT prompt that requests the model to create a *Solver* function to solve the problem and return the final answer.
+
+3. **Program Synthesis (PS) prompting:** Like PoT, it prompts the model to write a program to solve the problem. The prompt used is: *"Write a program that answers the following question: \{Problem\}"*.
+
+## Experiment Results
+
+The accuracy on all the problems with difficulty level-5 from different categories of the MATH dataset with different methods is shown below:
+
+![Result](img/result.png)
+
+We found that compared to basic prompting, which demonstrates the innate capabilities of GPT-4, utilizing Python within the context of PoT or PS strategy improved the overall accuracy by about 10%. This increase was mostly seen in categories involving more numerical manipulations, such as Counting & Probability and Number Theory, and in more complex categories like Intermediate Algebra and Precalculus.
+
+For categories like Algebra and Prealgebra, PoT and PS showed little improvement, and in some instances, even led to a decrease in accuracy. However, MathChat was able to enhance total accuracy by around 6% compared to PoT and PS, showing competitive performance across all categories. Remarkably, MathChat improved accuracy in the Algebra category by about 15% over other methods. Note that categories like Intermediate Algebra and Precalculus remained challenging for all methods, with only about 20% of problems solved accurately.
+
+The code for experiments can be found at this [repository](https://github.com/kevin666aa/FLAML/tree/gpt_math_solver/flaml/autogen/math).
+We now provide an implementation of MathChat using the interactive agents in FLAML. See this [notebook](https://github.com/microsoft/FLAML/blob/main/notebook/autogen_agentchat_MathChat.ipynb) for example usage.
+
+## Future Directions
+
+Despite MathChat's improvements over previous methods, the results show that complex math problem is still challenging for recent powerful LLMs, like GPT-4, even with help from external tools.
+
+Further work can be done to enhance this framework or math problem-solving in general:
+- Although enabling the model to use tools like Python can reduce calculation errors, LLMs are still prone to logic errors. Methods like self-consistency (Sample several solutions and take a major vote on the final answer), or self-verification (use another LLM instance to check whether an answer is correct) might improve the performance.
+- Sometimes, whether the LLM can solve the problem depends on the plan it uses. Some plans require less computation and logical reasoning, leaving less room for mistakes.
+- MathChat has the potential to be adapted into a copilot system, which could assist users with math problems. This system could allow users to be more involved in the problem-solving process, potentially enhancing learning.
+
+## For Further Reading
+
+* [Research paper of MathChat](https://arxiv.org/abs/2306.01337)
+* [Documentation about `flaml.autogen`](/docs/Use-Cases/Auto-Generation)
+
+*Are you working on applications that involve math problem-solving? Would you appreciate additional research or support on the application of LLM-based agents for math problem-solving? Please join our [Discord](https://discord.gg/Cppx2vSPVP) server for discussion.*
diff --git a/website/blog/authors.yml b/website/blog/authors.yml
index cd12283900..2aee7a5035 100644
--- a/website/blog/authors.yml
+++ b/website/blog/authors.yml
@@ -10,6 +10,12 @@ qingyunwu:
   url: https://qingyun-wu.github.io/
   image_url: https://github.com/qingyun-wu.png
 
+yiranwu:
+  name: Yiran Wu
+  title: PhD student at Pennsylvania State University
+  url: https://github.com/kevin666aa
+  image_url: https://github.com/kevin666aa.png
+
 jialeliu:
   name: Jiale Liu
   title: Undergraduate student at Xidian University
diff --git a/website/docs/Examples/AutoGen-OpenAI.md b/website/docs/Examples/AutoGen-OpenAI.md
index 97ce73df4b..cc7b208d7f 100644
--- a/website/docs/Examples/AutoGen-OpenAI.md
+++ b/website/docs/Examples/AutoGen-OpenAI.md
@@ -82,9 +82,9 @@ The tuning will be performed under the specified optimization budgets.
 Users can specify tuning data, optimization metric, optimization mode, evaluation function, search spaces etc.
 
 ```python
-from flaml import oai
+from flaml import autogen
 
-config, analysis = oai.Completion.tune(
+config, analysis = autogen.Completion.tune(
     data=tune_data,  # the data for tuning
     metric="success",  # the metric to optimize
     mode="max",  # the optimization mode
@@ -106,7 +106,7 @@ config, analysis = oai.Completion.tune(
 
 #### Output tuning results
 
-After the tuning, we can print out the optimized config and the result found by FLAML:
+After the tuning, we can print out the optimized config and the result:
 
 ```python
 print("optimized config", config)
@@ -118,17 +118,17 @@ print("best result on tuning data", analysis.best_result)
 We can apply the tuned config to the request for an instance:
 
 ```python
-response = oai.Completion.create(context=tune_data[1], **config)
+response = autogen.Completion.create(context=tune_data[1], **config)
 print(response)
-print(eval_with_generated_assertions(oai.Completion.extract_text(response), **tune_data[1]))
+print(eval_with_generated_assertions(autogen.Completion.extract_text(response), **tune_data[1]))
 ```
 
 #### Evaluate the success rate on the test data
 
-You can use flaml's `oai.Completion.test` to evaluate the performance of an entire dataset with the tuned config.
+You can use `autogen.Completion.test` to evaluate the performance of an entire dataset with the tuned config.
 
 ```python
-result = oai.Completion.test(test_data, **config)
+result = autogen.Completion.test(test_data, **config)
 print("performance on test data with the tuned config:", result)
 ```
 
diff --git a/website/docs/Getting-Started.md b/website/docs/Getting-Started.md
index 776c1f9cdf..ce86920045 100644
--- a/website/docs/Getting-Started.md
+++ b/website/docs/Getting-Started.md
@@ -25,10 +25,10 @@ There are several ways of using flaml:
 Maximize the utility out of the expensive LLMs such as ChatGPT and GPT-4, including:
 - A drop-in replacement of `openai.Completion` or `openai.ChatCompletion` with powerful functionalites like tuning, caching, templating, filtering. For example, you can optimize generations by LLM with your own tuning data, success metrics and budgets.
 ```python
-from flaml import oai
+from flaml import autogen
 
 # perform tuning
-config, analysis = oai.Completion.tune(
+config, analysis = autogen.Completion.tune(
     data=tune_data,
     metric="success",
     mode="max",
@@ -39,13 +39,13 @@ config, analysis = oai.Completion.tune(
 )
 
 # perform inference for a test instance
-response = oai.Completion.create(context=test_instance, **config)
+response = autogen.Completion.create(context=test_instance, **config)
 ```
 - LLM-driven intelligent agents which can perform tasks autonomously or with human feedback, including tasks that require using tools via code. For example,
 ```python
-assistant = AssistantAgent("assistant")
-user_proxy = UserProxyAgent("user_proxy")
-user_proxy.initiate_chat("Show me the YTD gain of 10 largest technology companies as of today.")
+assistant = autogen.AssistantAgent("assistant")
+user_proxy = autogen.UserProxyAgent("user_proxy")
+user_proxy.initiate_chat(assistant, message="Show me the YTD gain of 10 largest technology companies as of today.")
 ```
 
 #### [Task-oriented AutoML](/docs/Use-Cases/task-oriented-automl)
diff --git a/website/docs/Use-Cases/Auto-Generation.md b/website/docs/Use-Cases/Auto-Generation.md
index 4baa2281fd..5febbb03ef 100644
--- a/website/docs/Use-Cases/Auto-Generation.md
+++ b/website/docs/Use-Cases/Auto-Generation.md
@@ -15,22 +15,19 @@ The package is under active development with more features upcoming.
 
 [`flaml.autogen.agentchat`](/docs/reference/autogen/agentchat/agent) contains an experimental implementation of interactive agents which can adapt to human or simulated feedback. This subpackage is under active development.
 
-We have designed different classes of Agents that are capable of communicating with each other through the exchange of messages to collaboratively finish a task. An agent can communicate with other agents and perform actions. Different agents can differ in what actions they perform after receiving messages.
+We have designed a generic `ResponsiveAgent` class for Agents that are capable of communicating with each other through the exchange of messages to collaboratively finish a task. An agent can communicate with other agents and perform actions. Different agents can differ in what actions they perform after receiving messages. Two representative subclasses are `AssistantAgent` and `UserProxyAgent`.
 
-### `AssistantAgent`
+- `AssistantAgent`. Designed to act as an assistant by responding to user requests. It could write Python code (in a Python coding block) for a user to execute when a message (typically a description of a task that needs to be solved) is received. Under the hood, the Python code is written by LLM (e.g., GPT-4). It can also receive the execution results and suggest code with bug fix. Its behavior can be altered by passing a new system message. The LLM [inference](#enhanced-inference) configuration can be configured via `llm_config`.
+- `UserProxyAgent`. Serves as a proxy for the human user. Upon receiving a message, the UserProxyAgent will either solicit the human user's input or prepare an automatically generated reply. The chosen action depends on the settings of the `human_input_mode` and `max_consecutive_auto_reply` when the `UserProxyAgent` instance is constructed, and whether a human user input is available.
+By default, the automatically generated reply is crafted based on automatic code execution. The `UserProxyAgent` triggers code execution automatically when it detects an executable code block in the received message and no human user input is provided. Code execution can be disabled by setting `code_execution_config` to False. LLM-based response is disabled by default. It can be enabled by setting `llm_config` to a dict corresponding to the [inference](#enhanced-inference) configuration.
+When `llm_config` is set to a dict, `UserProxyAgent` can generate replies using an LLM when code execution is not performed.
 
-`AssistantAgent` is an Agent class designed to act as an assistant by responding to user requests. It could write Python code (in a Python coding block) for a user to execute when a message (typically a description of a task that needs to be solved) is received. Under the hood, the Python code is written by LLM (e.g., GPT-4).
-
-### `UserProxyAgent`
-`UserProxyAgent` is an Agent class that serves as a proxy for the human user. Upon receiving a message, the UserProxyAgent will either solicit the human user's input or prepare an automatically generated reply. The chosen action depends on the settings of the `human_input_mode` and `max_consecutive_auto_reply` when the `UserProxyAgent` instance is constructed, and whether a human user input is available.
-
-By default, the automatically generated reply is crafted based on automatic code execution. The `UserProxyAgent` triggers code execution automatically when it detects an executable code block in the received message and no human user input is provided. One can also easily extend it by overriding the `auto_reply` function of the `UserProxyAgent` to add or modify responses.
-For example, `AIUserProxyAgent` is a subclass of `UserProxyAgent` which can generate replies using an LLM when code execution is not performed. Code execution can be disabled by setting `code_execution_config` to False.
-This auto-reply capability allows for more autonomous user-agent communication while retaining the possibility of human intervention.
+The auto-reply capability of `ResponsiveAgent` allows for more autonomous multi-agent communication while retaining the possibility of human intervention.
+One can also easily extend it by overriding the `generate_reply` function of the `UserProxyAgent` to add or modify responses.
 
 Example usage of the agents to solve a task with code:
 ```python
-from flaml.autogen.agent import AssistantAgent, UserProxyAgent
+from flaml.autogen.agentchat import AssistantAgent, UserProxyAgent
 
 # create an AssistantAgent instance named "assistant"
 assistant = AssistantAgent(name="assistant")
@@ -39,9 +36,6 @@ assistant = AssistantAgent(name="assistant")
 user_proxy = UserProxyAgent(
     name="user_proxy",
     human_input_mode="NEVER",  # in this mode, the agent will never solicit human input but always auto reply
-    max_consecutive_auto_reply=10,  # the maximum number of consecutive auto replies
-    is_termination_msg=lambda x: x.get("content", "").rstrip().endswith("TERMINATE") or x.get("content", "").rstrip().endswith('"TERMINATE".'),  # the function to determine whether a message is a termination message
-    code_execution_config={"work_dir": "."},
 )
 
 # the assistant receives a message from the user, which contains the task description
@@ -53,7 +47,7 @@ user.initiate_chat(
 In the example above, we create an AssistantAgent named "assistant" to serve as the assistant and a UserProxyAgent named "user_proxy" to serve as a proxy for the human user.
 1. The assistant receives a message from the user_proxy, which contains the task description.
 2. The assistant then tries to write Python code to solve the task and sends the response to the user_proxy.
-3. Once the user_proxy receives a response from the assistant, it tries to reply by either soliciting human input or preparing an automatically generated reply. In this specific example, since `human_input_mode` is set to `"NEVER"`, the user_proxy will not solicit human input but prepare an automatically generated reply (auto reply). More specifically, the user_proxy executes the code and uses the result as the auto-reply.
+3. Once the user_proxy receives a response from the assistant, it tries to reply by either soliciting human input or preparing an automatically generated reply. In this specific example, since `human_input_mode` is set to `"NEVER"`, the user_proxy will not solicit human input but send an automatically generated reply (auto reply). More specifically, the user_proxy executes the code and uses the result as the auto-reply.
 4. The assistant then generates a further response for the user_proxy. The user_proxy can then decide whether to terminate the conversation. If not, steps 3 and 4 are repeated.
 
 Please find a visual illustration of how UserProxyAgent and AssistantAgent collaboratively solve the above task below:
@@ -70,57 +64,80 @@ To leverage [function calling capability of OpenAI's Chat Completions API](https
 
 Example usage of the agents to solve a task with function calling feature:
 ```python
-from flaml.autogen.agent import AssistantAgent, UserProxyAgent
+from flaml.autogen.agentchat import AssistantAgent, UserProxyAgent
 
 # put the descriptions of functions in config to be passed to OpenAI's API
-oai_config = {
+llm_config = {
     "model": "gpt-4-0613",
     "functions": [
         {
-            "name": "execute_code",
-            "description": "Receive a python code or shell script and return the execution result.",
+            "name": "python",
+            "description": "run cell in ipython and return the execution result.",
             "parameters": {
                 "type": "object",
                 "properties": {
-                    "code_type": {
+                    "cell": {
                         "type": "string",
-                        "description": "Code type, 'python' or 'sh'.",
-                    },
-                    "code": {
+                        "description": "Valid Python cell to execute.",
+                    }
+                },
+                "required": ["cell"],
+            },
+        },
+        {
+            "name": "sh",
+            "description": "run a shell script and return the execution result.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "script": {
                         "type": "string",
-                        "description": "Valid Python code to execute.",
+                        "description": "Valid shell script to execute.",
                     }
                 },
-                "required": ["code_type", "code"],
+                "required": ["script"],
             },
-        }
+        },
     ],
-    "function_call": "auto",
 }
 
 # create an AssistantAgent instance named "assistant"
-chatbot = AssistantAgent("assistant", config_list=config_list, **oai_config)
+chatbot = AssistantAgent("assistant", **llm_config)
 
-# create a UserProxyAgent instance named "user"
-user = UserProxyAgent(
-    "user",
+# create a UserProxyAgent instance named "user_proxy"
+user_proxy = UserProxyAgent(
+    "user_proxy",
     human_input_mode="NEVER",
-    code_execution_config={"work_dir": "coding"},
 )
 
-# define an `execute_code` function according to the function desription
-def execute_code(code_type, code):
-    # here we reuse the method in the user proxy agent
-    # in general, this is not necessary
-    return user.execute_code_blocks([(code_type, code)])
-
-# register the `execute_code` function
-user.register_function(function_map={"execute_code": execute_code})
+# define functions according to the function desription
+from IPython import get_ipython
+
+def exec_python(cell):
+    ipython = get_ipython()
+    result = ipython.run_cell(cell)
+    log = str(result.result)
+    if result.error_before_exec is not None:
+        log += f"\n{result.error_before_exec}"
+    if result.error_in_exec is not None:
+        log += f"\n{result.error_in_exec}"
+    return log
+
+def exec_sh(script):
+    return user_proxy.execute_code_blocks([("sh", script)])
+
+# register the functions
+user_proxy.register_function(
+    function_map={
+        "python": exec_python,
+        "sh": exec_sh,
+    }
+)
 
 # start the conversation
-user.initiate_chat(
-    assistant,
-    message="Draw a rocket and save to a file named 'rocket.svg'",
+user_proxy.initiate_chat(
+    chatbot,
+    message="Draw two agents chatting with each other with an example dialog.",
 )
 ```
 
@@ -131,8 +148,6 @@ user.initiate_chat(
 
 * [Solve Tasks Requiring Web Info](https://github.com/microsoft/FLAML/blob/main/notebook/autogen_agentchat_web_info.ipynb)
 
-* [Using MathChat to Solve Math Problems](https://github.com/microsoft/FLAML/blob/main/notebook/autogen_agentchat_MathChat.ipynb)
-
 * [Use Provided Tools as Functions](https://github.com/microsoft/FLAML/blob/main/notebook/autogen_agentchat_function_call.ipynb)
 
 * [Automated Task Solving with Coding & Planning Agents](https://github.com/microsoft/FLAML/blob/main/notebook/autogen_agentchat_planning.ipynb)
@@ -141,8 +156,8 @@ user.initiate_chat(
 
 ## Enhanced Inference
 
-One can use [`flaml.oai.Completion.create`](/docs/reference/autogen/oai/completion#create) to perform inference.
-There are a number of benefits of using `flaml.oai.Completion.create` to perform inference.
+One can use [`flaml.autogen.Completion.create`](/docs/reference/autogen/oai/completion#create) to perform inference.
+There are a number of benefits of using `flaml.autogen.Completion.create` to perform inference.
 
 ### Tune Inference Parameters
 
@@ -222,12 +237,12 @@ The optimization budget refers to the total budget allowed in the tuning process
 
 #### Perform tuning
 
-Now, you can use [`flaml.oai.Completion.tune`](/docs/reference/autogen/oai/completion#tune) for tuning. For example,
+Now, you can use [`flaml.autogen.Completion.tune`](/docs/reference/autogen/oai/completion#tune) for tuning. For example,
 
 ```python
-from flaml import oai
+from flaml import autogen
 
-config, analysis = oai.Completion.tune(
+config, analysis = autogen.Completion.tune(
     data=tune_data,
     metric="success",
     mode="max",
@@ -250,12 +265,12 @@ The tuend config can be used to perform inference.
 
 ### API unification
 
-`flaml.oai.Completion.create` is compatible with both `openai.Completion.create` and `openai.ChatCompletion.create`, and both OpenAI API and Azure OpenAI API. So models such as "text-davinci-003", "gpt-3.5-turbo" and "gpt-4" can share a common API.
-When chat models are used and `prompt` is given as the input to `flaml.oai.Completion.create`, the prompt will be automatically converted into `messages` to fit the chat completion API requirement. One advantage is that one can experiment with both chat and non-chat models for the same prompt in a unified API.
+`flaml.autogen.Completion.create` is compatible with both `openai.Completion.create` and `openai.ChatCompletion.create`, and both OpenAI API and Azure OpenAI API. So models such as "text-davinci-003", "gpt-3.5-turbo" and "gpt-4" can share a common API.
+When chat models are used and `prompt` is given as the input to `flaml.autogen.Completion.create`, the prompt will be automatically converted into `messages` to fit the chat completion API requirement. One advantage is that one can experiment with both chat and non-chat models for the same prompt in a unified API.
 
 For local LLMs, one can spin up an endpoint using a package like [simple_ai_server](https://github.com/lhenault/simpleAI) and [FastChat](https://github.com/lm-sys/FastChat), and then use the same API to send a request. See [here](/blog/2023/07/14/Local-LLMs) for examples on how to make inference with local LLMs.
 
-When only working with the chat-based models, `flaml.oai.ChatCompletion` can be used. It also does automatic conversion from prompt to messages, if prompt is provided instead of messages.
+When only working with the chat-based models, `flaml.autogen.ChatCompletion` can be used. It also does automatic conversion from prompt to messages, if prompt is provided instead of messages.
 
 ### Caching
 
@@ -265,12 +280,12 @@ API call results are cached locally and reused when the same request is issued.
 
 #### Runtime error
 
-It is easy to hit error when calling OpenAI APIs, due to connection, rate limit, or timeout. Some of the errors are transient. `flaml.oai.Completion.create` deals with the transient errors and retries automatically. Initial request timeout, retry timeout and retry time interval can be configured via `request_timeout`, `retry_timeout` and `flaml.oai.Completion.retry_time`.
+It is easy to hit error when calling OpenAI APIs, due to connection, rate limit, or timeout. Some of the errors are transient. `flaml.autogen.Completion.create` deals with the transient errors and retries automatically. Initial request timeout, retry timeout and retry time interval can be configured via `request_timeout`, `retry_timeout` and `flaml.autogen.Completion.retry_time`.
 
 Moreover, one can pass a list of configurations of different models/endpoints to mitigate the rate limits. For example,
 
 ```python
-response = oai.Completion.create(
+response = autogen.Completion.create(
     config_list=[
         {
             "model": "gpt-4",
@@ -306,7 +321,7 @@ Another type of error is that the returned response does not satisfy a requireme
 
 ```python
 def valid_json_filter(context, config, response):
-    for text in oai.Completion.extract_text(response):
+    for text in autogen.Completion.extract_text(response):
         try:
             json.loads(text)
             return True
@@ -314,7 +329,7 @@ def valid_json_filter(context, config, response):
             pass
     return False
 
-response = oai.Completion.create(
+response = autogen.Completion.create(
     config_list=[{"model": "text-ada-001"}, {"model": "gpt-3.5-turbo"}, {"model": "text-davinci-003"}],
     prompt="How to construct a json request to Bing API to search for 'latest AI news'? Return the JSON request.",
     filter_func=valid_json_filter,
@@ -330,7 +345,7 @@ The example above will try to use text-ada-001, gpt-3.5-turbo, and text-davinci-
 If the provided prompt or message is a template, it will be automatically materialized with a given context. For example,
 
 ```python
-response = oai.Completion.create(
+response = autogen.Completion.create(
     context={"problem": "How many positive integers, not exceeding 100, are multiples of 2 or 3 but not 4?"},
     prompt="{problem} Solve the problem carefully.",
     **config
@@ -363,11 +378,11 @@ context = {
     "external_info_0": "Problem 1: ...",
 }
 
-response = oai.ChatCompletion.create(context, messages=messages, **config)
+response = autogen.ChatCompletion.create(context, messages=messages, **config)
 messages.append(
     {
         "role": "assistant",
-        "content": oai.ChatCompletion.extract_text(response)[0]
+        "content": autogen.ChatCompletion.extract_text(response)[0]
     }
 )
 messages.append(
@@ -382,26 +397,26 @@ context.append(
         "external_info_1": "Theorem 1: ...",
     }
 )
-response = oai.ChatCompletion.create(context, messages=messages, **config)
+response = autogen.ChatCompletion.create(context, messages=messages, **config)
 ```
 
 ### Logging (Experimental)
 
-When debugging or diagnosing an LLM-based system, it is often convenient to log the API calls and analyze them. `flaml.oai.Completion` and `flaml.oai.ChatCompletion` offer an easy way to collect the API call histories. For example, to log the chat histories, simply run:
+When debugging or diagnosing an LLM-based system, it is often convenient to log the API calls and analyze them. `flaml.autogen.Completion` and `flaml.autogen.ChatCompletion` offer an easy way to collect the API call histories. For example, to log the chat histories, simply run:
 ```python
-flaml.oai.ChatCompletion.start_logging()
+flaml.autogen.ChatCompletion.start_logging()
 ```
 The API calls made after this will be automatically logged. They can be retrieved at any time by:
 ```python
-flaml.oai.ChatCompletion.logged_history
+flaml.autogen.ChatCompletion.logged_history
 ```
 To stop logging, use
 ```python
-flaml.oai.ChatCompletion.stop_logging()
+flaml.autogen.ChatCompletion.stop_logging()
 ```
 If one would like to append the history to an existing dict, pass the dict like:
 ```python
-flaml.oai.ChatCompletion.start_logging(history_dict=existing_history_dict)
+flaml.autogen.ChatCompletion.start_logging(history_dict=existing_history_dict)
 ```
 By default, the counter of API calls will be reset at `start_logging()`. If no reset is desired, set `reset_counter=False`.
 


Test patch:
diff --git a/test/autogen/extensions/__init__.py b/test/autogen/agentchat/extensions/__init__.py
similarity index 100%
rename from test/autogen/extensions/__init__.py
rename to test/autogen/agentchat/extensions/__init__.py
diff --git a/test/autogen/extensions/tsp.py b/test/autogen/agentchat/extensions/tsp.py
similarity index 100%
rename from test/autogen/extensions/tsp.py
rename to test/autogen/agentchat/extensions/tsp.py
diff --git a/test/autogen/extensions/tsp_api.py b/test/autogen/agentchat/extensions/tsp_api.py
similarity index 100%
rename from test/autogen/extensions/tsp_api.py
rename to test/autogen/agentchat/extensions/tsp_api.py
diff --git a/test/autogen/test_assistant_agent.py b/test/autogen/agentchat/test_assistant_agent.py
similarity index 86%
rename from test/autogen/test_assistant_agent.py
rename to test/autogen/agentchat/test_assistant_agent.py
index 2b194c2663..a13ed535af 100644
--- a/test/autogen/test_assistant_agent.py
+++ b/test/autogen/agentchat/test_assistant_agent.py
@@ -1,10 +1,10 @@
 import os
 import sys
 import pytest
-from flaml import oai
+from flaml import autogen
 from flaml.autogen.agentchat import AssistantAgent, UserProxyAgent
 
-KEY_LOC = "test/autogen"
+KEY_LOC = "notebook"
 OAI_CONFIG_LIST = "OAI_CONFIG_LIST"
 here = os.path.abspath(os.path.dirname(__file__))
 
@@ -20,16 +20,16 @@ def test_ai_user_proxy_agent():
         return
 
     conversations = {}
-    oai.ChatCompletion.start_logging(conversations)
+    autogen.ChatCompletion.start_logging(conversations)
 
-    config_list = oai.config_list_from_json(
+    config_list = autogen.config_list_from_json(
         OAI_CONFIG_LIST,
         file_location=KEY_LOC,
     )
     assistant = AssistantAgent(
         "assistant",
         system_message="You are a helpful assistant.",
-        oai_config={
+        llm_config={
             "request_timeout": 600,
             "seed": 42,
             "config_list": config_list,
@@ -41,7 +41,7 @@ def test_ai_user_proxy_agent():
         human_input_mode="NEVER",
         max_consecutive_auto_reply=2,
         code_execution_config=False,
-        oai_config={
+        llm_config={
             "config_list": config_list,
         },
         # In the system message the "user" always refers to ther other agent.
@@ -62,7 +62,7 @@ def test_gpt35(human_input_mode="NEVER", max_consecutive_auto_reply=5):
         import openai
     except ImportError:
         return
-    config_list = oai.config_list_from_json(
+    config_list = autogen.config_list_from_json(
         OAI_CONFIG_LIST,
         file_location=KEY_LOC,
         filter_dict={
@@ -77,7 +77,7 @@ def test_gpt35(human_input_mode="NEVER", max_consecutive_auto_reply=5):
     )
     assistant = AssistantAgent(
         "coding_agent",
-        oai_config={
+        llm_config={
             # "request_timeout": 600,
             "seed": 42,
             "config_list": config_list,
@@ -115,12 +115,12 @@ def test_create_execute_script(human_input_mode="NEVER", max_consecutive_auto_re
     except ImportError:
         return
 
-    config_list = oai.config_list_from_json(OAI_CONFIG_LIST, file_location=KEY_LOC)
+    config_list = autogen.config_list_from_json(OAI_CONFIG_LIST, file_location=KEY_LOC)
     conversations = {}
-    oai.ChatCompletion.start_logging(conversations)
+    autogen.ChatCompletion.start_logging(conversations)
     assistant = AssistantAgent(
         "assistant",
-        oai_config={
+        llm_config={
             "request_timeout": 600,
             "seed": 42,
             "config_list": config_list,
@@ -145,10 +145,10 @@ def test_create_execute_script(human_input_mode="NEVER", max_consecutive_auto_re
 ```""",
     )
     print(conversations)
-    oai.ChatCompletion.start_logging(compact=False)
+    autogen.ChatCompletion.start_logging(compact=False)
     user.send("""Execute temp.py""", assistant)
-    print(oai.ChatCompletion.logged_history)
-    oai.ChatCompletion.stop_logging()
+    print(autogen.ChatCompletion.logged_history)
+    autogen.ChatCompletion.stop_logging()
 
 
 def test_tsp(human_input_mode="NEVER", max_consecutive_auto_reply=10):
@@ -157,7 +157,7 @@ def test_tsp(human_input_mode="NEVER", max_consecutive_auto_reply=10):
     except ImportError:
         return
 
-    config_list = oai.config_list_from_json(
+    config_list = autogen.config_list_from_json(
         OAI_CONFIG_LIST,
         file_location=KEY_LOC,
         filter_dict={
@@ -179,19 +179,17 @@ def __init__(self, *args, **kwargs):
         def generate_init_message(self, question) -> str:
             return self._prompt.format(question=question)
 
-    oai.ChatCompletion.start_logging()
-    assistant = AssistantAgent("assistant", oai_config={"temperature": 0, "config_list": config_list})
+    autogen.ChatCompletion.start_logging()
+    assistant = AssistantAgent("assistant", llm_config={"temperature": 0, "config_list": config_list})
     user = TSPUserProxyAgent(
         "user",
         code_execution_config={"work_dir": here},
         human_input_mode=human_input_mode,
         max_consecutive_auto_reply=max_consecutive_auto_reply,
     )
-    # agent.receive(prompt.format(question=hard_questions[0]), user)
-    # agent.receive(prompt.format(question=hard_questions[1]), user)
     user.initiate_chat(assistant, question=hard_questions[2])
-    print(oai.ChatCompletion.logged_history)
-    oai.ChatCompletion.stop_logging()
+    print(autogen.ChatCompletion.logged_history)
+    autogen.ChatCompletion.stop_logging()
 
 
 if __name__ == "__main__":
diff --git a/test/autogen/test_math_user_proxy_agent.py b/test/autogen/agentchat/test_math_user_proxy_agent.py
similarity index 94%
rename from test/autogen/test_math_user_proxy_agent.py
rename to test/autogen/agentchat/test_math_user_proxy_agent.py
index 569e9195f1..83c800a39f 100644
--- a/test/autogen/test_math_user_proxy_agent.py
+++ b/test/autogen/agentchat/test_math_user_proxy_agent.py
@@ -1,14 +1,12 @@
-from flaml import oai
+import pytest
+import sys
+from flaml import autogen
 from flaml.autogen.agentchat.contrib.math_user_proxy_agent import (
     MathUserProxyAgent,
     _remove_print,
     _add_print_to_last_line,
 )
-import pytest
-import sys
-
-KEY_LOC = "test/autogen"
-OAI_CONFIG_LIST = "OAI_CONFIG_LIST"
+from test_assistant_agent import KEY_LOC, OAI_CONFIG_LIST
 
 
 @pytest.mark.skipif(
@@ -24,9 +22,9 @@ def test_math_user_proxy_agent():
     from flaml.autogen.agentchat.assistant_agent import AssistantAgent
 
     conversations = {}
-    oai.ChatCompletion.start_logging(conversations)
+    autogen.ChatCompletion.start_logging(conversations)
 
-    config_list = oai.config_list_from_json(
+    config_list = autogen.config_list_from_json(
         OAI_CONFIG_LIST,
         file_location=KEY_LOC,
         filter_dict={
@@ -36,7 +34,7 @@ def test_math_user_proxy_agent():
     assistant = AssistantAgent(
         "assistant",
         system_message="You are a helpful assistant.",
-        oai_config={
+        llm_config={
             "request_timeout": 600,
             "seed": 42,
             "config_list": config_list,
diff --git a/test/autogen/agentchat/test_optiguide.py b/test/autogen/agentchat/test_optiguide.py
new file mode 100644
index 0000000000..8b6be251f8
--- /dev/null
+++ b/test/autogen/agentchat/test_optiguide.py
@@ -0,0 +1,65 @@
+import pytest
+import sys
+import requests  # for loading the example source code
+from flaml import autogen
+from test_assistant_agent import KEY_LOC, OAI_CONFIG_LIST
+
+
+@pytest.mark.skipif(
+    sys.platform in ["darwin", "win32"],
+    reason="do not run on MacOS or windows",
+)
+def test_optiguide():
+    try:
+        import openai
+    except ImportError:
+        return
+
+    from flaml.autogen.agentchat import UserProxyAgent
+    from flaml.autogen.agentchat.contrib.opti_guide import OptiGuideAgent
+
+    conversations = {}
+    autogen.ChatCompletion.start_logging(conversations)
+
+    config_list = autogen.config_list_from_json(
+        OAI_CONFIG_LIST,
+        file_location=KEY_LOC,
+        filter_dict={
+            "model": ["gpt-4", "gpt4", "gpt-4-32k", "gpt-4-32k-0314"],
+        },
+    )
+    code_url = "https://www.beibinli.com/docs/optiguide/coffee.py"
+    code = requests.get(code_url).text
+
+    assistant = OptiGuideAgent(
+        "optiguide",
+        source_code=code,
+        example_qa="",
+        llm_config={
+            "request_timeout": 600,
+            "seed": 42,
+            "config_list": config_list,
+        },
+    )
+
+    user_proxy = UserProxyAgent(name="user_proxy", human_input_mode="NEVER", max_consecutive_auto_reply=0)
+    user_proxy.send("What if we prohibit shipping from supplier 1 to roastery 2?", assistant)
+    user_proxy.send(
+        "What is the impact of supplier1 being able to supply only half the quantity at present?", assistant
+    )
+    user_proxy.send("What if Roastery 1 is exclusively for Cafe 3", assistant)
+
+    # test danger case
+    assistant.reset()
+    assistant.debug_times = 1
+    assistant._shield.generate_reply = lambda **_: "DANGER"
+    user_proxy.send("What's the weather today", assistant)
+    assert (
+        assistant._debug_times_left == 0
+        and assistant.last_message(user_proxy)["content"] == "Sorry. I cannot answer your question."
+    )
+    # print(conversations)
+
+
+if __name__ == "__main__":
+    test_optiguide()
diff --git a/test/autogen/test_responsive_agent.py b/test/autogen/agentchat/test_responsive_agent.py
similarity index 75%
rename from test/autogen/test_responsive_agent.py
rename to test/autogen/agentchat/test_responsive_agent.py
index 5710a5a27a..6656f6c90b 100644
--- a/test/autogen/test_responsive_agent.py
+++ b/test/autogen/agentchat/test_responsive_agent.py
@@ -19,10 +19,11 @@ def test_responsive_agent(monkeypatch):
     )  # receive a dict
 
     # receive dict without openai fields to be printed, such as "content", 'function_call'. There should be no error raised.
-    pre_len = len(dummy_agent_1.oai_conversations["dummy_agent_2"])
-    dummy_agent_1.receive({"message": "hello"}, dummy_agent_2)
+    pre_len = len(dummy_agent_1.chat_messages["dummy_agent_2"])
+    with pytest.raises(ValueError):
+        dummy_agent_1.receive({"message": "hello"}, dummy_agent_2)
     assert pre_len == len(
-        dummy_agent_1.oai_conversations["dummy_agent_2"]
+        dummy_agent_1.chat_messages["dummy_agent_2"]
     ), "When the message is not an valid openai message, it should not be appended to the oai conversation."
 
     monkeypatch.setattr(sys, "stdin", StringIO("exit"))
@@ -36,14 +37,18 @@ def test_responsive_agent(monkeypatch):
     )  # send a dict
 
     # send dict with no openai fields
-    pre_len = len(dummy_agent_1.oai_conversations["dummy_agent_2"])
+    pre_len = len(dummy_agent_1.chat_messages["dummy_agent_2"])
     with pytest.raises(ValueError):
         dummy_agent_1.send({"message": "hello"}, dummy_agent_2)
 
     assert pre_len == len(
-        dummy_agent_1.oai_conversations["dummy_agent_2"]
+        dummy_agent_1.chat_messages["dummy_agent_2"]
     ), "When the message is not a valid openai message, it should not be appended to the oai conversation."
 
+    # update system message
+    dummy_agent_1.update_system_message("new system message")
+    assert dummy_agent_1._oai_system_message[0]["content"] == "new system message"
+
 
 if __name__ == "__main__":
     test_responsive_agent(pytest.monkeypatch)
diff --git a/test/autogen/tsp_prompt.txt b/test/autogen/agentchat/tsp_prompt.txt
similarity index 100%
rename from test/autogen/tsp_prompt.txt
rename to test/autogen/agentchat/tsp_prompt.txt
diff --git a/test/autogen/oai/test_completion.py b/test/autogen/oai/test_completion.py
index 97d6042564..33fa5ed364 100644
--- a/test/autogen/oai/test_completion.py
+++ b/test/autogen/oai/test_completion.py
@@ -5,7 +5,7 @@
 from functools import partial
 import os
 import json
-from flaml import oai
+from flaml import autogen
 from flaml.autogen.code_utils import (
     eval_function_completions,
     generate_assertions,
@@ -16,19 +16,19 @@
 )
 from flaml.autogen.math_utils import eval_math_responses, solve_problem
 
-KEY_LOC = "test/autogen"
+KEY_LOC = "notebook"
 OAI_CONFIG_LIST = "OAI_CONFIG_LIST"
 here = os.path.abspath(os.path.dirname(__file__))
 
 
 def yes_or_no_filter(context, response, **_):
     return context.get("yes_or_no_choice", False) is False or any(
-        text in ["Yes.", "No."] for text in oai.Completion.extract_text(response)
+        text in ["Yes.", "No."] for text in autogen.Completion.extract_text(response)
     )
 
 
 def valid_json_filter(response, **_):
-    for text in oai.Completion.extract_text(response):
+    for text in autogen.Completion.extract_text(response):
         try:
             json.loads(text)
             return True
@@ -43,47 +43,47 @@ def test_filter():
     except ImportError as exc:
         print(exc)
         return
-    response = oai.Completion.create(
+    response = autogen.Completion.create(
         context={"yes_or_no_choice": True},
         config_list=[{"model": "text-ada-001"}, {"model": "gpt-3.5-turbo"}, {"model": "text-davinci-003"}],
         prompt="Is 37 a prime number? Please answer 'Yes.' or 'No.'",
         filter_func=yes_or_no_filter,
     )
     assert (
-        oai.Completion.extract_text(response)[0] in ["Yes.", "No."]
+        autogen.Completion.extract_text(response)[0] in ["Yes.", "No."]
         or not response["pass_filter"]
         and response["config_id"] == 2
     )
-    response = oai.Completion.create(
+    response = autogen.Completion.create(
         context={"yes_or_no_choice": False},
         config_list=[{"model": "text-ada-001"}, {"model": "gpt-3.5-turbo"}, {"model": "text-davinci-003"}],
         prompt="Is 37 a prime number?",
         filter_func=yes_or_no_filter,
     )
     assert response["model"] == "text-ada-001"
-    response = oai.Completion.create(
+    response = autogen.Completion.create(
         config_list=[{"model": "text-ada-001"}, {"model": "gpt-3.5-turbo"}, {"model": "text-davinci-003"}],
         prompt="How to construct a json request to Bing API to search for 'latest AI news'? Return the JSON request.",
         filter_func=valid_json_filter,
     )
     assert response["config_id"] == 2 or response["pass_filter"], "the response must pass filter unless all fail"
-    assert not response["pass_filter"] or json.loads(oai.Completion.extract_text(response)[0])
+    assert not response["pass_filter"] or json.loads(autogen.Completion.extract_text(response)[0])
 
 
 def test_chatcompletion():
-    params = oai.ChatCompletion._construct_params(
+    params = autogen.ChatCompletion._construct_params(
         context=None,
         config={"model": "unknown"},
         prompt="hi",
     )
     assert "messages" in params
-    params = oai.Completion._construct_params(
+    params = autogen.Completion._construct_params(
         context=None,
         config={"model": "unknown"},
         prompt="hi",
     )
     assert "messages" not in params
-    params = oai.Completion._construct_params(
+    params = autogen.Completion._construct_params(
         context=None,
         config={"model": "gpt-4"},
         prompt="hi",
@@ -97,8 +97,8 @@ def test_multi_model():
     except ImportError as exc:
         print(exc)
         return
-    response = oai.Completion.create(
-        config_list=oai.config_list_gpt4_gpt35(KEY_LOC),
+    response = autogen.Completion.create(
+        config_list=autogen.config_list_gpt4_gpt35(KEY_LOC),
         prompt="Hi",
     )
     print(response)
@@ -111,7 +111,7 @@ def test_improve():
     except ImportError as exc:
         print(exc)
         return
-    config_list = oai.config_list_openai_aoai(KEY_LOC)
+    config_list = autogen.config_list_openai_aoai(KEY_LOC)
     improved, _ = improve_function(
         "flaml/autogen/math_utils.py",
         "solve_problem",
@@ -144,12 +144,12 @@ def test_nocontext():
     except ImportError as exc:
         print(exc)
         return
-    response = oai.Completion.create(
+    response = autogen.Completion.create(
         model="text-ada-001", prompt="1+1=", max_tokens=1, use_cache=False, request_timeout=10
     )
     print(response)
     code, _ = generate_code(
-        config_list=oai.config_list_from_json(
+        config_list=autogen.config_list_from_json(
             OAI_CONFIG_LIST,
             file_location=KEY_LOC,
             filter_dict={
@@ -175,7 +175,7 @@ def test_nocontext():
     )
     print(code)
 
-    solution, cost = solve_problem("1+1=", config_list=oai.config_list_gpt4_gpt35(KEY_LOC))
+    solution, cost = solve_problem("1+1=", config_list=autogen.config_list_gpt4_gpt35(KEY_LOC))
     print(solution, cost)
 
 
@@ -184,7 +184,7 @@ def test_nocontext():
     reason="do not run on windows",
 )
 def test_humaneval(num_samples=1):
-    gpt35_config_list = oai.config_list_from_json(
+    gpt35_config_list = autogen.config_list_from_json(
         env_or_file="OAI_CONFIG_LIST",
         filter_dict={
             "model": {
@@ -221,17 +221,17 @@ def test_humaneval(num_samples=1):
         }
         for x in range(n_tune_data, len(data))
     ]
-    oai.Completion.clear_cache(cache_path_root="{here}/cache")
-    oai.Completion.set_cache(seed)
+    autogen.Completion.clear_cache(cache_path_root="{here}/cache")
+    autogen.Completion.set_cache(seed)
     try:
         import openai
         import diskcache
     except ImportError as exc:
         print(exc)
         return
-    oai.Completion.clear_cache(400)
+    autogen.Completion.clear_cache(400)
     # no error should be raised
-    response = oai.Completion.create(
+    response = autogen.Completion.create(
         context=test_data[0],
         config_list=[{"model": "gpt-3.5-turbo"}],
         prompt="",
@@ -241,7 +241,7 @@ def test_humaneval(num_samples=1):
     )
     # assert response == -1
     # a minimal tuning example
-    config, _ = oai.Completion.tune(
+    config, _ = autogen.Completion.tune(
         data=tune_data,
         metric="success",
         mode="max",
@@ -249,9 +249,9 @@ def test_humaneval(num_samples=1):
         n=1,
         prompt="{definition}",
     )
-    response = oai.Completion.create(context=test_data[0], **config)
+    response = autogen.Completion.create(context=test_data[0], **config)
     # a minimal tuning example for tuning chat completion models using the Completion class
-    config, _ = oai.Completion.tune(
+    config, _ = autogen.Completion.tune(
         data=tune_data,
         metric="succeed_assertions",
         mode="max",
@@ -260,10 +260,10 @@ def test_humaneval(num_samples=1):
         model="text-davinci-003",
         prompt="{definition}",
     )
-    response = oai.Completion.create(context=test_data[0], **config)
+    response = autogen.Completion.create(context=test_data[0], **config)
     # a minimal tuning example for tuning chat completion models using the ChatCompletion class
-    config_list = oai.config_list_openai_aoai(KEY_LOC)
-    config, _ = oai.ChatCompletion.tune(
+    config_list = autogen.config_list_openai_aoai(KEY_LOC)
+    config, _ = autogen.ChatCompletion.tune(
         data=tune_data,
         metric="expected_success",
         mode="max",
@@ -272,7 +272,7 @@ def test_humaneval(num_samples=1):
         messages=[{"role": "user", "content": "{definition}"}],
         config_list=config_list,
     )
-    response = oai.ChatCompletion.create(context=test_data[0], config_list=config_list, **config)
+    response = autogen.ChatCompletion.create(context=test_data[0], config_list=config_list, **config)
     print(response)
     from openai.error import RateLimitError
 
@@ -289,7 +289,7 @@ def test_humaneval(num_samples=1):
     assert selected == 0
     print(eval_function_completions([code], **tune_data[1]))
     # a more comprehensive tuning example
-    config2, analysis = oai.Completion.tune(
+    config2, analysis = autogen.Completion.tune(
         data=tune_data,
         metric="success",
         mode="max",
@@ -310,12 +310,12 @@ def test_humaneval(num_samples=1):
     print(config2)
     print(analysis.best_result)
     print(test_data[0])
-    response = oai.Completion.create(context=test_data[0], **config2)
+    response = autogen.Completion.create(context=test_data[0], **config2)
     print(response)
-    oai.Completion.data = test_data[:num_samples]
-    result = oai.Completion._eval(analysis.best_config, prune=False, eval_only=True)
+    autogen.Completion.data = test_data[:num_samples]
+    result = autogen.Completion._eval(analysis.best_config, prune=False, eval_only=True)
     print("result without pruning", result)
-    result = oai.Completion.test(test_data[:num_samples], **config2)
+    result = autogen.Completion.test(test_data[:num_samples], **config2)
     print(result)
     try:
         code, cost, selected = implement(
@@ -376,7 +376,7 @@ def test_math(num_samples=-1):
         % data["problem"]
     ]
 
-    oai.Completion.set_cache(seed)
+    autogen.Completion.set_cache(seed)
     vanilla_config = {
         "model": "text-davinci-003",
         "temperature": 1,
@@ -386,8 +386,8 @@ def test_math(num_samples=-1):
         "stop": "###",
     }
     test_data_sample = test_data[0:3]
-    result = oai.Completion.test(test_data_sample, eval_math_responses, **vanilla_config)
-    result = oai.Completion.test(
+    result = autogen.Completion.test(test_data_sample, eval_math_responses, **vanilla_config)
+    result = autogen.Completion.test(
         test_data_sample,
         eval_math_responses,
         agg_method="median",
@@ -400,13 +400,13 @@ def my_median(results):
     def my_average(results):
         return np.mean(results)
 
-    result = oai.Completion.test(
+    result = autogen.Completion.test(
         test_data_sample,
         eval_math_responses,
         agg_method=my_median,
         **vanilla_config,
     )
-    result = oai.Completion.test(
+    result = autogen.Completion.test(
         test_data_sample,
         eval_math_responses,
         agg_method={
@@ -420,7 +420,7 @@ def my_average(results):
 
     print(result)
 
-    config, _ = oai.Completion.tune(
+    config, _ = autogen.Completion.tune(
         data=tune_data,  # the data for tuning
         metric="expected_success",  # the metric to optimize
         mode="max",  # the optimization mode
@@ -433,7 +433,7 @@ def my_average(results):
         stop="###",  # the stop sequence
     )
     print("tuned config", config)
-    result = oai.Completion.test(test_data_sample, config_list=oai.config_list_openai_aoai(KEY_LOC), **config)
+    result = autogen.Completion.test(test_data_sample, config_list=autogen.config_list_openai_aoai(KEY_LOC), **config)
     print("result from tuned config:", result)
     print("empty responses", eval_math_responses([], None))
 
@@ -441,7 +441,7 @@ def my_average(results):
 if __name__ == "__main__":
     import openai
 
-    config_list = oai.config_list_openai_aoai(KEY_LOC)
+    config_list = autogen.config_list_openai_aoai(KEY_LOC)
     assert len(config_list) >= 3, config_list
     openai.api_key = os.environ["OPENAI_API_KEY"]
 
diff --git a/test/autogen/oai/test_utils.py b/test/autogen/oai/test_utils.py
index 2789e5cba8..17ed33d7f6 100644
--- a/test/autogen/oai/test_utils.py
+++ b/test/autogen/oai/test_utils.py
@@ -1,22 +1,21 @@
 import json
 import os
-from flaml import oai
-
-KEY_LOC = "test/autogen"
+from flaml import autogen
+from test_completion import KEY_LOC, OAI_CONFIG_LIST
 
 
 def test_config_list_from_json():
-    config_list = oai.config_list_gpt4_gpt35(key_file_path=KEY_LOC)
+    config_list = autogen.config_list_gpt4_gpt35(key_file_path=KEY_LOC)
     json_file = os.path.join(KEY_LOC, "config_list_test.json")
     with open(json_file, "w") as f:
         json.dump(config_list, f, indent=4)
-    config_list_1 = oai.config_list_from_json(json_file)
+    config_list_1 = autogen.config_list_from_json(json_file)
     assert config_list == config_list_1
     os.environ["config_list_test"] = json.dumps(config_list)
-    config_list_2 = oai.config_list_from_json("config_list_test")
+    config_list_2 = autogen.config_list_from_json("config_list_test")
     assert config_list == config_list_2
-    config_list_3 = oai.config_list_from_json(
-        "OAI_CONFIG_LIST", file_location=KEY_LOC, filter_dict={"model": ["gpt4", "gpt-4-32k"]}
+    config_list_3 = autogen.config_list_from_json(
+        OAI_CONFIG_LIST, file_location=KEY_LOC, filter_dict={"model": ["gpt4", "gpt-4-32k"]}
     )
     assert all(config.get("model") in ["gpt4", "gpt-4-32k"] for config in config_list_3)
     del os.environ["config_list_test"]
@@ -24,7 +23,7 @@ def test_config_list_from_json():
 
 
 def test_config_list_openai_aoai():
-    config_list = oai.config_list_openai_aoai(key_file_path=KEY_LOC)
+    config_list = autogen.config_list_openai_aoai(key_file_path=KEY_LOC)
     assert all(config.get("api_type") in [None, "open_ai", "azure"] for config in config_list)
 
 
diff --git a/test/autogen/test_code.py b/test/autogen/test_code.py
index af7652b8d3..7f7938990c 100644
--- a/test/autogen/test_code.py
+++ b/test/autogen/test_code.py
@@ -59,12 +59,16 @@ def test_execute_code():
         import docker
     except ImportError as exc:
         print(exc)
-        return
-    exitcode, msg, image = execute_code("print('hello world')", filename="tmp/codetest.py")
-    assert exitcode == 0 and msg == b"hello world\n", msg
+        docker = None
+    exit_code, msg, image = execute_code("print('hello world')", filename="tmp/codetest.py")
+    assert exit_code == 0 and msg == "hello world\n", msg
     # read a file
     print(execute_code("with open('tmp/codetest.py', 'r') as f: a=f.read()"))
     # create a file
+    exit_code, msg, image = execute_code(
+        "with open('tmp/codetest.py', 'w') as f: f.write('b=1')", work_dir=f"{here}/my_tmp", filename="tmp2/codetest.py"
+    )
+    assert exit_code and 'File "tmp2/codetest.py"' in msg, msg
     print(execute_code("with open('tmp/codetest.py', 'w') as f: f.write('b=1')", work_dir=f"{here}/my_tmp"))
     # execute code in a file
     print(execute_code(filename="tmp/codetest.py"))
@@ -72,16 +76,17 @@ def test_execute_code():
     # execute code for assertion error
     exit_code, msg, image = execute_code("assert 1==2")
     assert exit_code, msg
+    assert 'File ""' in msg
     # execute code which takes a long time
     exit_code, error, image = execute_code("import time; time.sleep(2)", timeout=1)
-    assert exit_code and error.decode() == "Timeout"
-    assert isinstance(image, str)
+    assert exit_code and error == "Timeout"
+    assert isinstance(image, str) or docker is None or os.path.exists("/.dockerenv")
 
 
 def test_execute_code_no_docker():
     exit_code, error, image = execute_code("import time; time.sleep(2)", timeout=1, use_docker=False)
     if sys.platform != "win32":
-        assert exit_code and error.decode() == "Timeout"
+        assert exit_code and error == "Timeout"
     assert image is None
 
 
diff --git a/test/autogen/test_function_call.py b/test/autogen/test_function_call.py
index ebce5f018f..91ce251569 100644
--- a/test/autogen/test_function_call.py
+++ b/test/autogen/test_function_call.py
@@ -4,15 +4,15 @@
     openai = None
 import pytest
 import json
-from flaml import oai
+from flaml import autogen
 from flaml.autogen.math_utils import eval_math_responses
 
-KEY_LOC = "test/autogen"
+KEY_LOC = "notebook"
 
 
 @pytest.mark.skipif(openai is None, reason="openai not installed")
 def test_eval_math_responses():
-    config_list = oai.config_list_from_models(
+    config_list = autogen.config_list_from_models(
         KEY_LOC, exclude="aoai", model_list=["gpt-4-0613", "gpt-3.5-turbo-0613", "gpt-3.5-turbo-16k"]
     )
     functions = [
@@ -36,7 +36,7 @@ def test_eval_math_responses():
             },
         },
     ]
-    response = oai.ChatCompletion.create(
+    response = autogen.ChatCompletion.create(
         config_list=config_list,
         messages=[
             {
@@ -47,7 +47,7 @@ def test_eval_math_responses():
         functions=functions,
     )
     print(response)
-    responses = oai.ChatCompletion.extract_text_or_function_call(response)
+    responses = autogen.ChatCompletion.extract_text_or_function_call(response)
     print(responses[0])
     function_call = responses[0]["function_call"]
     name, arguments = function_call["name"], json.loads(function_call["arguments"])