Fix third party imports using optional import block in agent contrib …

…files (#585)

* Move third party non default module imports to optional block in autogen/agentchat/contrib files

* Fix failing test_retrievechat tests

* Fix typo in require_optional_import decorator

* Move third party non default module imports to optional block in autogen/agentchat/contrib/vectordb files

* Move third party non default module imports to optional block in autogen/agentchat/contrib/graph_rag files

* Fix failing falkor graph rag test

* Move third party non default module imports to optional block in autogen/agentchat/contrib/captainagent files

* Move third party non default module imports to optional block in autogen/agentchat/contrib/capabilities files

* Move third party non default module imports to optional block in autogen/agentchat/contrib/agent_eval files

* Fix imports in missed files

autogen/agentchat/contrib/agent_eval/agent_eval.py

@@ -6,12 +6,12 @@
 # SPDX-License-Identifier: MIT
 from typing import Literal, Optional, Union
 
-import autogen
-from autogen.agentchat.contrib.agent_eval.criterion import Criterion
-from autogen.agentchat.contrib.agent_eval.critic_agent import CriticAgent
-from autogen.agentchat.contrib.agent_eval.quantifier_agent import QuantifierAgent
-from autogen.agentchat.contrib.agent_eval.subcritic_agent import SubCriticAgent
-from autogen.agentchat.contrib.agent_eval.task import Task
+from .... import GroupChat, GroupChatManager, UserProxyAgent
+from .criterion import Criterion
+from .critic_agent import CriticAgent
+from .quantifier_agent import QuantifierAgent
+from .subcritic_agent import SubCriticAgent
+from .task import Task
 
 
 def generate_criteria(
@@ -38,7 +38,7 @@ def generate_criteria(
         llm_config=llm_config,
     )
 
-    critic_user = autogen.UserProxyAgent(
+    critic_user = UserProxyAgent(
         name="critic_user",
         max_consecutive_auto_reply=0,  # terminate without auto-reply
         human_input_mode="NEVER",
@@ -53,10 +53,8 @@ def generate_criteria(
         )
         agents.append(subcritic)
 
-    groupchat = autogen.GroupChat(
-        agents=agents, messages=[], max_round=max_round, speaker_selection_method="round_robin"
-    )
-    critic_manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=llm_config)
+    groupchat = GroupChat(agents=agents, messages=[], max_round=max_round, speaker_selection_method="round_robin")
+    critic_manager = GroupChatManager(groupchat=groupchat, llm_config=llm_config)
 
     critic_user.initiate_chat(critic_manager, message=task.get_sys_message())
     criteria = critic_user.last_message()
@@ -90,7 +88,7 @@ def quantify_criteria(
         llm_config=llm_config,
     )
 
-    quantifier_user = autogen.UserProxyAgent(
+    quantifier_user = UserProxyAgent(
         name="quantifier_user",
         max_consecutive_auto_reply=0,  # terminate without auto-reply
         human_input_mode="NEVER",

autogen/agentchat/contrib/agent_eval/critic_agent.py

@@ -6,7 +6,7 @@
 # SPDX-License-Identifier: MIT
 from typing import Optional
 
-from autogen.agentchat.conversable_agent import ConversableAgent
+from ...conversable_agent import ConversableAgent
 
 
 class CriticAgent(ConversableAgent):

autogen/agentchat/contrib/agent_eval/quantifier_agent.py

@@ -6,7 +6,7 @@
 # SPDX-License-Identifier: MIT
 from typing import Optional
 
-from autogen.agentchat.conversable_agent import ConversableAgent
+from ...conversable_agent import ConversableAgent
 
 
 class QuantifierAgent(ConversableAgent):

autogen/agentchat/contrib/agent_eval/subcritic_agent.py

@@ -6,7 +6,7 @@
 # SPDX-License-Identifier: MIT
 from typing import Optional
 
-from autogen.agentchat.conversable_agent import ConversableAgent
+from ...conversable_agent import ConversableAgent
 
 
 class SubCriticAgent(ConversableAgent):

autogen/agentchat/contrib/agent_optimizer.py

@@ -8,8 +8,8 @@
 import json
 from typing import Optional
 
-import autogen
-from autogen.code_utils import execute_code
+from ... import OpenAIWrapper, filter_config
+from ...code_utils import execute_code
 
 ADD_FUNC = {
     "type": "function",
@@ -209,10 +209,8 @@ def __init__(
             raise ValueError(
                 "When using OpenAI or Azure OpenAI endpoints, specify a non-empty 'model' either in 'llm_config' or in each config of 'config_list'."
             )
-        self.llm_config["config_list"] = autogen.filter_config(
-            llm_config["config_list"], {"model": [self.optimizer_model]}
-        )
-        self._client = autogen.OpenAIWrapper(**self.llm_config)
+        self.llm_config["config_list"] = filter_config(llm_config["config_list"], {"model": [self.optimizer_model]})
+        self._client = OpenAIWrapper(**self.llm_config)
 
     def record_one_conversation(self, conversation_history: list[dict], is_satisfied: bool = None):
         """Record one conversation history.

autogen/agentchat/contrib/capabilities/agent_capability.py

@@ -4,7 +4,7 @@
 #
 # Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
 # SPDX-License-Identifier: MIT
-from autogen.agentchat.assistant_agent import ConversableAgent
+from ...assistant_agent import ConversableAgent
 
 
 class AgentCapability:

autogen/agentchat/contrib/capabilities/teachability.py

@@ -8,14 +8,15 @@
 import pickle
 from typing import Optional, Union
 
-import chromadb
-from chromadb.config import Settings
-
-from autogen.agentchat.assistant_agent import ConversableAgent
-from autogen.agentchat.contrib.capabilities.agent_capability import AgentCapability
-from autogen.agentchat.contrib.text_analyzer_agent import TextAnalyzerAgent
-
 from ....formatting_utils import colored
+from ....import_utils import optional_import_block, require_optional_import
+from ...assistant_agent import ConversableAgent
+from ..text_analyzer_agent import TextAnalyzerAgent
+from .agent_capability import AgentCapability
+
+with optional_import_block():
+    import chromadb
+    from chromadb.config import Settings
 
 
 class Teachability(AgentCapability):
@@ -238,6 +239,7 @@ def _analyze(self, text_to_analyze: Union[dict, str], analysis_instructions: Uni
         return self.teachable_agent.last_message(self.analyzer)["content"]
 
 
+@require_optional_import("chromadb", "teachable")
 class MemoStore:
     """Provides memory storage and retrieval for a teachable agent, using a vector database.
     Each DB entry (called a memo) is a pair of strings: an input text and an output text.

autogen/agentchat/contrib/capabilities/tools_capability.py

@@ -2,8 +2,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from autogen.agentchat import ConversableAgent
-from autogen.tools import Tool
+from ....agentchat import ConversableAgent
+from ....tools import Tool
 
 
 class ToolsCapability:

autogen/agentchat/contrib/capabilities/transforms.py

@@ -11,10 +11,9 @@
 import tiktoken
 from termcolor import colored
 
-from autogen import token_count_utils
-from autogen.cache import AbstractCache, Cache
-from autogen.types import MessageContentType
-
+from .... import token_count_utils
+from ....cache import AbstractCache, Cache
+from ....types import MessageContentType
 from . import transforms_util
 from .text_compressors import LLMLingua, TextCompressor
 

autogen/agentchat/contrib/capabilities/transforms_util.py

@@ -7,10 +7,10 @@
 from collections.abc import Hashable
 from typing import Any, Optional
 
-from autogen import token_count_utils
-from autogen.cache.abstract_cache_base import AbstractCache
-from autogen.oai.openai_utils import filter_config
-from autogen.types import MessageContentType
+from .... import token_count_utils
+from ....cache.abstract_cache_base import AbstractCache
+from ....oai.openai_utils import filter_config
+from ....types import MessageContentType
 
 
 def cache_key(content: MessageContentType, *args: Hashable) -> str:

autogen/agentchat/contrib/capabilities/vision_capability.py

@@ -7,16 +7,16 @@
 import copy
 from typing import Callable, Optional, Union
 
-from autogen.agentchat.assistant_agent import ConversableAgent
-from autogen.agentchat.contrib.capabilities.agent_capability import AgentCapability
-from autogen.agentchat.contrib.img_utils import (
+from ....code_utils import content_str
+from ....oai.client import OpenAIWrapper
+from ...assistant_agent import ConversableAgent
+from ..img_utils import (
     convert_base64_to_data_uri,
     get_image_data,
     get_pil_image,
     gpt4v_formatter,
 )
-from autogen.code_utils import content_str
-from autogen.oai.client import OpenAIWrapper
+from .agent_capability import AgentCapability
 
 DEFAULT_DESCRIPTION_PROMPT = (
     "Write a detailed caption for this image. "

autogen/agentchat/contrib/captainagent/agent_builder.py

@@ -15,7 +15,8 @@
 
 from termcolor import colored
 
-import autogen
+from .... import AssistantAgent, ConversableAgent, OpenAIWrapper, UserProxyAgent, config_list_from_json
+from ....code_utils import CODE_BLOCK_PATTERN
 
 __all__ = ["AgentBuilder"]
 
@@ -37,7 +38,7 @@ def _config_check(config: dict):
 
 
 def _retrieve_json(text):
-    match = re.findall(autogen.code_utils.CODE_BLOCK_PATTERN, text, flags=re.DOTALL)
+    match = re.findall(CODE_BLOCK_PATTERN, text, flags=re.DOTALL)
     if not match:
         return text
     code_blocks = []
@@ -203,15 +204,15 @@ def __init__(
             builder_filter_dict.update({"model": builder_model})
         if len(builder_model_tags) != 0:
             builder_filter_dict.update({"tags": builder_model_tags})
-        builder_config_list = autogen.config_list_from_json(
+        builder_config_list = config_list_from_json(
             config_file_or_env, file_location=config_file_location, filter_dict=builder_filter_dict
         )
         if len(builder_config_list) == 0:
             raise RuntimeError(
                 f"Fail to initialize build manager: {builder_model}{builder_model_tags} does not exist in {config_file_or_env}. "
                 f'If you want to change this model, please specify the "builder_model" in the constructor.'
             )
-        self.builder_model = autogen.OpenAIWrapper(config_list=builder_config_list)
+        self.builder_model = OpenAIWrapper(config_list=builder_config_list)
 
         self.agent_model = agent_model if isinstance(agent_model, list) else [agent_model]
         self.agent_model_tags = agent_model_tags
@@ -222,7 +223,7 @@ def __init__(
         self.agent_configs: list[dict] = []
         self.open_ports: list[str] = []
         self.agent_procs: dict[str, tuple[sp.Popen, str]] = {}
-        self.agent_procs_assign: dict[str, tuple[autogen.ConversableAgent, str]] = {}
+        self.agent_procs_assign: dict[str, tuple[ConversableAgent, str]] = {}
         self.cached_configs: dict = {}
 
         self.max_agents = max_agents
@@ -239,7 +240,7 @@ def _create_agent(
         member_name: list[str],
         llm_config: dict,
         use_oai_assistant: Optional[bool] = False,
-    ) -> autogen.AssistantAgent:
+    ) -> AssistantAgent:
         """Create a group chat participant agent.
 
         If the agent rely on an open-source model, this function will automatically set up an endpoint for that agent.
@@ -274,7 +275,7 @@ def _create_agent(
             filter_dict.update({"model": model_name_or_hf_repo})
         if len(model_tags) > 0:
             filter_dict.update({"tags": model_tags})
-        config_list = autogen.config_list_from_json(
+        config_list = config_list_from_json(
             self.config_file_or_env, file_location=self.config_file_location, filter_dict=filter_dict
         )
         if len(config_list) == 0:
@@ -287,7 +288,7 @@ def _create_agent(
         current_config = llm_config.copy()
         current_config.update({"config_list": config_list})
         if use_oai_assistant:
-            from autogen.agentchat.contrib.gpt_assistant_agent import GPTAssistantAgent
+            from ..gpt_assistant_agent import GPTAssistantAgent
 
             agent = GPTAssistantAgent(
                 name=agent_name,
@@ -302,14 +303,14 @@ def _create_agent(
                     "\nThe group also include a Computer_terminal to help you run the python and shell code."
                 )
 
-            model_class = autogen.AssistantAgent
+            model_class = AssistantAgent
             if agent_path:
                 module_path, model_class_name = agent_path.replace("/", ".").rsplit(".", 1)
                 module = importlib.import_module(module_path)
                 model_class = getattr(module, model_class_name)
-                if not issubclass(model_class, autogen.ConversableAgent):
+                if not issubclass(model_class, ConversableAgent):
                     logger.error(f"{model_class} is not a ConversableAgent. Use AssistantAgent as default")
-                    model_class = autogen.AssistantAgent
+                    model_class = AssistantAgent
 
             additional_config = {
                 k: v
@@ -365,10 +366,10 @@ def build(
         coding: Optional[bool] = None,
         code_execution_config: Optional[dict] = None,
         use_oai_assistant: Optional[bool] = False,
-        user_proxy: Optional[autogen.ConversableAgent] = None,
+        user_proxy: Optional[ConversableAgent] = None,
         max_agents: Optional[int] = None,
         **kwargs,
-    ) -> tuple[list[autogen.ConversableAgent], dict]:
+    ) -> tuple[list[ConversableAgent], dict]:
         """Auto build agents based on the building task.
 
         Args:
@@ -496,9 +497,9 @@ def build_from_library(
         code_execution_config: Optional[dict] = None,
         use_oai_assistant: Optional[bool] = False,
         embedding_model: Optional[str] = "all-mpnet-base-v2",
-        user_proxy: Optional[autogen.ConversableAgent] = None,
+        user_proxy: Optional[ConversableAgent] = None,
         **kwargs,
-    ) -> tuple[list[autogen.ConversableAgent], dict]:
+    ) -> tuple[list[ConversableAgent], dict]:
         """Build agents from a library.
         The library is a list of agent configs, which contains the name and system_message for each agent.
         We use a build manager to decide what agent in that library should be involved to the task.
@@ -655,8 +656,8 @@ def build_from_library(
         return self._build_agents(use_oai_assistant, user_proxy=user_proxy, **kwargs)
 
     def _build_agents(
-        self, use_oai_assistant: Optional[bool] = False, user_proxy: Optional[autogen.ConversableAgent] = None, **kwargs
-    ) -> tuple[list[autogen.ConversableAgent], dict]:
+        self, use_oai_assistant: Optional[bool] = False, user_proxy: Optional[ConversableAgent] = None, **kwargs
+    ) -> tuple[list[ConversableAgent], dict]:
         """Build agents with generated configs.
 
         Args:
@@ -687,7 +688,7 @@ def _build_agents(
         if coding is True:
             print("Adding user console proxy...", flush=True)
             if user_proxy is None:
-                user_proxy = autogen.UserProxyAgent(
+                user_proxy = UserProxyAgent(
                     name="Computer_terminal",
                     is_termination_msg=lambda x: x == "TERMINATE" or x == "TERMINATE.",
                     code_execution_config=code_execution_config,
@@ -722,7 +723,7 @@ def load(
         config_json: Optional[str] = None,
         use_oai_assistant: Optional[bool] = False,
         **kwargs,
-    ) -> tuple[list[autogen.ConversableAgent], dict]:
+    ) -> tuple[list[ConversableAgent], dict]:
         """Load building configs and call the build function to complete building without calling online LLMs' api.
 
         Args:

autogen/agentchat/contrib/captainagent/captainagent.py

@@ -8,10 +8,8 @@
 
 from termcolor import colored
 
-import autogen
-from autogen import UserProxyAgent
-from autogen.agentchat.conversable_agent import ConversableAgent
-
+from .... import GroupChat, GroupChatManager, UserProxyAgent
+from ...conversable_agent import ConversableAgent
 from .agent_builder import AgentBuilder
 from .tool_retriever import ToolBuilder, format_ag2_tool, get_full_tool_description
 
@@ -465,13 +463,13 @@ def _run_autobuild(self, group_name: str, execution_task: str, building_task: st
 
         self.build_times += 1
         # start nested chat
-        nested_group_chat = autogen.GroupChat(
+        nested_group_chat = GroupChat(
             agents=agent_list,
             messages=[],
             allow_repeat_speaker=agent_list[:-1] if agent_configs["coding"] is True else agent_list,
             **self._nested_config["group_chat_config"],
         )
-        manager = autogen.GroupChatManager(
+        manager = GroupChatManager(
             groupchat=nested_group_chat,
             llm_config=self._nested_config["group_chat_llm_config"],
         )

autogen/agentchat/contrib/captainagent/tool_retriever.py

@@ -16,15 +16,18 @@
 from textwrap import dedent, indent
 from typing import Optional, Union
 
-import pandas as pd
-from sentence_transformers import SentenceTransformer, util
-
 from .... import AssistantAgent, UserProxyAgent
 from ....coding import CodeExecutor, CodeExtractor, LocalCommandLineCodeExecutor, MarkdownCodeExtractor
 from ....coding.base import CodeBlock, CodeResult
+from ....import_utils import optional_import_block, require_optional_import
 from ....tools import Tool, get_function_schema, load_basemodels_if_needed
 
+with optional_import_block():
+    import pandas as pd
+    from sentence_transformers import SentenceTransformer, util
+
 
+@require_optional_import(["pandas", "sentence_transformers"], "retrievechat")
 class ToolBuilder:
     TOOL_PROMPT_DEFAULT = """\n## Functions
 You have access to the following functions. They can be accessed from the module called 'functions' by their function names.