Refactor GPTAssistantAgent

autogen/agentchat/contrib/gpt_assistant_agent.py

@@ -7,6 +7,7 @@
 from autogen import OpenAIWrapper
 from autogen.agentchat.agent import Agent
 from autogen.agentchat.assistant_agent import ConversableAgent
+from autogen.agentchat.assistant_agent import AssistantAgent
 from typing import Dict, Optional, Union, List, Tuple, Any
 
 logger = logging.getLogger(__name__)
@@ -21,45 +22,76 @@ class GPTAssistantAgent(ConversableAgent):
     def __init__(
         self,
         name="GPT Assistant",
-        instructions: Optional[str] = "You are a helpful GPT Assistant.",
+        instructions: Optional[str] = None,
         llm_config: Optional[Union[Dict, bool]] = None,
+        overwrite_instructions: bool = False,
     ):
         """
         Args:
             name (str): name of the agent.
             instructions (str): instructions for the OpenAI assistant configuration.
+            When instructions is not None, the system message of the agent will be
+            set to the provided instructions and used in the assistant run, irrespective
+            of the overwrite_instructions flag. But when instructions is None,
+            and the assistant does not exist, the system message will be set to
+            AssistantAgent.DEFAULT_SYSTEM_MESSAGE. If the assistant exists, the
+            system message will be set to the existing assistant instructions.
             llm_config (dict or False): llm inference configuration.
+                - assistant_id: ID of the assistant to use. If None, a new assistant will be created.
                 - model: Model to use for the assistant (gpt-4-1106-preview, gpt-3.5-turbo-1106).
                 - check_every_ms: check thread run status interval
                 - tools: Give Assistants access to OpenAI-hosted tools like Code Interpreter and Knowledge Retrieval,
                         or build your own tools using Function calling. ref https://platform.openai.com/docs/assistants/tools
                 - file_ids: files used by retrieval in run
+            overwrite_instructions (bool): whether to overwrite the instructions of an existing assistant.
         """
-        super().__init__(
-            name=name,
-            system_message=instructions,
-            human_input_mode="NEVER",
-            llm_config=llm_config,
-        )
-
         # Use AutoGen OpenAIWrapper to create a client
-        oai_wrapper = OpenAIWrapper(**self.llm_config)
+        oai_wrapper = OpenAIWrapper(**llm_config)
         if len(oai_wrapper._clients) > 1:
             logger.warning("GPT Assistant only supports one OpenAI client. Using the first client in the list.")
         self._openai_client = oai_wrapper._clients[0]
-
         openai_assistant_id = llm_config.get("assistant_id", None)
         if openai_assistant_id is None:
             # create a new assistant
+            if instructions is None:
+                logger.warning(
+                    "No instructions were provided for new assistant. Using default instructions from AssistantAgent.DEFAULT_SYSTEM_MESSAGE."
+                )
+                instructions = AssistantAgent.DEFAULT_SYSTEM_MESSAGE
             self._openai_assistant = self._openai_client.beta.assistants.create(
                 name=name,
                 instructions=instructions,
-                tools=self.llm_config.get("tools", []),
-                model=self.llm_config.get("model", "gpt-4-1106-preview"),
+                tools=llm_config.get("tools", []),
+                model=llm_config.get("model", "gpt-4-1106-preview"),
             )
         else:
             # retrieve an existing assistant
             self._openai_assistant = self._openai_client.beta.assistants.retrieve(openai_assistant_id)
+            # if no instructions are provided, set the instructions to the existing instructions
+            if instructions is None:
+                logger.warning(
+                    "No instructions were provided for given assistant. Using existing instructions from assistant API."
+                )
+                instructions = self.get_assistant_instructions()
+            elif overwrite_instructions is True:
+                logger.warning(
+                    "overwrite_instructions is True. Provided instructions will be used and will modify the assistant in the API"
+                )
+                self._openai_assistant = self._openai_client.beta.assistants.update(
+                    assistant_id=openai_assistant_id,
+                    instructions=instructions,
+                )
+            else:
+                logger.warning(
+                    "overwrite_instructions is False. Provided instructions will be used without permanently modifying the assistant in the API."
+                )
+
+        super().__init__(
+            name=name,
+            system_message=instructions,
+            human_input_mode="NEVER",
+            llm_config=llm_config,
+        )
 
         # lazly create thread
         self._openai_threads = {}
@@ -107,6 +139,8 @@ def _invoke_assistant(
         run = self._openai_client.beta.threads.runs.create(
             thread_id=assistant_thread.id,
             assistant_id=self._openai_assistant.id,
+            # pass the latest system message as instructions
+            instructions=self.system_message,
         )
 
         run_response_messages = self._get_run_response(assistant_thread, run)
@@ -300,3 +334,16 @@ def pretty_print_thread(self, thread):
     def oai_threads(self) -> Dict[Agent, Any]:
         """Return the threads of the agent."""
         return self._openai_threads
+
+    @property
+    def assistant_id(self):
+        """Return the assistant id"""
+        return self._openai_assistant.id
+
+    def get_assistant_instructions(self):
+        """Return the assistant instructions from OAI assistant API"""
+        return self._openai_assistant.instructions
+
+    def delete_assistant(self):
+        """Delete the assistant from OAI assistant API"""
+        self._openai_client.beta.assistants.delete(self.assistant_id)

test/agentchat/contrib/test_gpt_assistant.py

@@ -38,9 +38,10 @@ def test_gpt_assistant_chat():
         "description": "This is an API endpoint allowing users (analysts) to input question about GitHub in text format to retrieve the realted and structured data.",
     }
 
+    config_list = autogen.config_list_from_json(OAI_CONFIG_LIST, file_location=KEY_LOC)
     analyst = GPTAssistantAgent(
         name="Open_Source_Project_Analyst",
-        llm_config={"tools": [{"type": "function", "function": ossinsight_api_schema}]},
+        llm_config={"tools": [{"type": "function", "function": ossinsight_api_schema}], "config_list": config_list},
         instructions="Hello, Open Source Project Analyst. You'll conduct comprehensive evaluations of open source projects or organizations on the GitHub platform",
     )
     analyst.register_function(
@@ -62,5 +63,102 @@ def test_gpt_assistant_chat():
     assert len(analyst._openai_threads) == 0
 
 
+def test_get_assistant_instructions():
+    """
+    Test function to create a new GPTAssistantAgent, set its instructions, retrieve the instructions,
+    and assert that the retrieved instructions match the set instructions.
+    """
+
+    config_list = autogen.config_list_from_json(OAI_CONFIG_LIST, file_location=KEY_LOC)
+    assistant = GPTAssistantAgent(
+        "assistant",
+        instructions="This is a test",
+        llm_config={
+            "config_list": config_list,
+        },
+    )
+
+    instruction_match = assistant.get_assistant_instructions() == "This is a test"
+    assistant.delete_assistant()
+
+    assert instruction_match is True
+
+
+def test_gpt_assistant_instructions_overwrite():
+    """
+    Test that the instructions of a GPTAssistantAgent can be overwritten or not depending on the value of the
+    `overwrite_instructions` parameter when creating a new assistant with the same ID.
+
+    Steps:
+    1. Create a new GPTAssistantAgent with some instructions.
+    2. Get the ID of the assistant.
+    3. Create a new GPTAssistantAgent with the same ID but different instructions and `overwrite_instructions=True`.
+    4. Check that the instructions of the assistant have been overwritten with the new ones.
+    """
+
+    instructions1 = "This is a test #1"
+    instructions2 = "This is a test #2"
+
+    config_list = autogen.config_list_from_json(OAI_CONFIG_LIST, file_location=KEY_LOC)
+    assistant = GPTAssistantAgent(
+        "assistant",
+        instructions=instructions1,
+        llm_config={
+            "config_list": config_list,
+        },
+    )
+
+    assistant_id = assistant.assistant_id
+    assistant = GPTAssistantAgent(
+        "assistant",
+        instructions=instructions2,
+        llm_config={
+            "config_list": config_list,
+            "assistant_id": assistant_id,
+        },
+        overwrite_instructions=True,
+    )
+
+    instruction_match = assistant.get_assistant_instructions() == instructions2
+    assistant.delete_assistant()
+
+    assert instruction_match is True
+
+
+def test_gpt_assistant_existing_no_instructions():
+    """
+    Test function to check if the GPTAssistantAgent can retrieve instructions for an existing assistant
+    even if the assistant was created with no instructions initially.
+    """
+    instructions = "This is a test #1"
+
+    config_list = autogen.config_list_from_json(OAI_CONFIG_LIST, file_location=KEY_LOC)
+    assistant = GPTAssistantAgent(
+        "assistant",
+        instructions=instructions,
+        llm_config={
+            "config_list": config_list,
+        },
+    )
+
+    assistant_id = assistant.assistant_id
+
+    # create a new assistant with the same ID but no instructions
+    assistant = GPTAssistantAgent(
+        "assistant",
+        llm_config={
+            "config_list": config_list,
+            "assistant_id": assistant_id,
+        },
+    )
+
+    instruction_match = assistant.get_assistant_instructions() == instructions
+    assistant.delete_assistant()
+    assert instruction_match is True
+
+
 if __name__ == "__main__":
     test_gpt_assistant_chat()
+    test_get_assistant_instructions()
+    test_gpt_assistant_instructions_overwrite()
+    test_gpt_assistant_existing_no_instructions()