AgentBuilder

def set_builder_model(self, model: str):
    self.builder_model = model

def set_agent_model(self, model: str):
    self.agent_model = model

def clear_agent(self, agent_name: str, recycle_endpoint: Optional[bool] = True):
    """Clear a specific agent by name.

    Args:
        agent_name: the name of agent.
        recycle_endpoint: trigger for recycle the endpoint server. If true, the endpoint will be recycled
            when there is no agent depending on.
    """
    _, server_id = self.agent_procs_assign[agent_name]
    del self.agent_procs_assign[agent_name]
    if recycle_endpoint:
        if server_id == self.online_server_name:
            return
        else:
            for _, iter_sid in self.agent_procs_assign.values():
                if server_id == iter_sid:
                    return
            self.agent_procs[server_id][0].terminate()
            self.open_ports.append(server_id.split("_")[-1])
    print(colored(f"Agent {agent_name} has been cleared.", "yellow"), flush=True)

def clear_all_agents(self, recycle_endpoint: Optional[bool] = True):
    """Clear all cached agents."""
    for agent_name in [agent_name for agent_name in self.agent_procs_assign]:
        self.clear_agent(agent_name, recycle_endpoint)
    print(colored("All agents have been cleared.", "yellow"), flush=True)

def build(
    self,
    building_task: str,
    default_llm_config: dict[str, Any],
    coding: Optional[bool] = None,
    code_execution_config: Optional[dict[str, Any]] = None,
    use_oai_assistant: Optional[bool] = False,
    user_proxy: Optional[ConversableAgent] = None,
    max_agents: Optional[int] = None,
    **kwargs: Any,
) -> tuple[list[ConversableAgent], dict[str, Any]]:
    """Auto build agents based on the building task.

    Args:
        building_task: instruction that helps build manager (gpt-4) to decide what agent should be built.
        default_llm_config: specific configs for LLM (e.g., config_list, seed, temperature, ...).
        coding: use to identify if the user proxy (a code interpreter) should be added.
        code_execution_config: specific configs for user proxy (e.g., last_n_messages, work_dir, ...).
        use_oai_assistant: use OpenAI assistant api instead of self-constructed agent.
        user_proxy: user proxy's class that can be used to replace the default user proxy.
        max_agents (Optional[int], default=None): Maximum number of agents to create for the task. If None, uses the value from self.max_agents.
        **kwargs (Any): Additional arguments to pass to _build_agents.
            - agent_configs: Optional list of predefined agent configurations to use.

    Returns:
        agent_list: a list of agents.
        cached_configs: cached configs.
    """
    if code_execution_config is None:
        code_execution_config = {
            "last_n_messages": 1,
            "work_dir": "groupchat",
            "use_docker": False,
            "timeout": 10,
        }

    if max_agents is None:
        max_agents = self.max_agents

    agent_configs = kwargs.get("agent_configs", [])
    self.building_task = building_task

    print(colored("==> Generating agents...", "green"), flush=True)
    resp_agent_name = (
        self.builder_model.create(
            messages=[
                {
                    "role": "user",
                    "content": self.AGENT_NAME_PROMPT.format(task=building_task, max_agents=max_agents),
                }
            ]
        )
        .choices[0]
        .message.content
    )
    agent_name_list = [agent_name.strip().replace(" ", "_") for agent_name in resp_agent_name.split(",")]
    print(f"{agent_name_list} are generated.", flush=True)

    print(colored("==> Generating system message...", "green"), flush=True)
    agent_sys_msg_list = []
    for name in agent_name_list:
        print(f"Preparing system message for {name}", flush=True)
        resp_agent_sys_msg = (
            self.builder_model.create(
                messages=[
                    {
                        "role": "user",
                        "content": self.AGENT_SYS_MSG_PROMPT.format(
                            task=building_task,
                            position=name,
                            default_sys_msg=self.DEFAULT_DESCRIPTION,
                        ),
                    }
                ]
            )
            .choices[0]
            .message.content
        )
        agent_sys_msg_list.append(resp_agent_sys_msg)

    print(colored("==> Generating description...", "green"), flush=True)
    agent_description_list = []
    for name, sys_msg in list(zip(agent_name_list, agent_sys_msg_list)):
        print(f"Preparing description for {name}", flush=True)
        resp_agent_description = (
            self.builder_model.create(
                messages=[
                    {
                        "role": "user",
                        "content": self.AGENT_DESCRIPTION_PROMPT.format(position=name, sys_msg=sys_msg),
                    }
                ]
            )
            .choices[0]
            .message.content
        )
        agent_description_list.append(resp_agent_description)

    for name, sys_msg, description in list(zip(agent_name_list, agent_sys_msg_list, agent_description_list)):
        agent_configs.append({
            "name": name,
            "model": self.agent_model,
            "tags": self.agent_model_tags,
            "system_message": sys_msg,
            "description": description,
        })

    if coding is None:
        resp = (
            self.builder_model.create(
                messages=[{"role": "user", "content": self.CODING_PROMPT.format(task=building_task)}]
            )
            .choices[0]
            .message.content
        )
        coding = resp == "YES"

    self.cached_configs.update({
        "building_task": building_task,
        "agent_configs": agent_configs,
        "coding": coding,
        "default_llm_config": default_llm_config,
        "code_execution_config": code_execution_config,
    })
    _config_check(self.cached_configs)
    return self._build_agents(use_oai_assistant, user_proxy=user_proxy, **kwargs)

def build_from_library(
    self,
    building_task: str,
    library_path_or_json: str,
    default_llm_config: dict[str, Any],
    top_k: int = 3,
    coding: Optional[bool] = None,
    code_execution_config: Optional[dict[str, Any]] = None,
    use_oai_assistant: Optional[bool] = False,
    embedding_model: Optional[str] = "all-mpnet-base-v2",
    user_proxy: Optional[ConversableAgent] = None,
    **kwargs: Any,
) -> tuple[list[ConversableAgent], dict[str, Any]]:
    """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.

    Args:
        building_task: instruction that helps build manager (gpt-4) to decide what agent should be built.
        library_path_or_json: path or JSON string config of agent library.
        default_llm_config: specific configs for LLM (e.g., config_list, seed, temperature, ...).
        top_k: number of results to return.
        coding: use to identify if the user proxy (a code interpreter) should be added.
        code_execution_config: specific configs for user proxy (e.g., last_n_messages, work_dir, ...).
        use_oai_assistant: use OpenAI assistant api instead of self-constructed agent.
        embedding_model: a Sentence-Transformers model use for embedding similarity to select agents from library.
            As reference, chromadb use "all-mpnet-base-v2" as default.
        user_proxy: user proxy's class that can be used to replace the default user proxy.
        **kwargs: Additional arguments to pass to _build_agents.

    Returns:
        agent_list: a list of agents.
        cached_configs: cached configs.
    """
    import sqlite3

    # Some system will have an unexpected sqlite3 version.
    # Check if the user has installed pysqlite3.
    if int(sqlite3.version.split(".")[0]) < 3:
        try:
            __import__("pysqlite3")
            import sys

            sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
        except Exception as e:
            raise e
    import chromadb
    from chromadb.utils import embedding_functions

    if code_execution_config is None:
        code_execution_config = {
            "last_n_messages": 1,
            "work_dir": "groupchat",
            "use_docker": False,
            "timeout": 120,
        }

    try:
        agent_library = json.loads(library_path_or_json)
    except json.decoder.JSONDecodeError:
        with open(library_path_or_json) as f:
            agent_library = json.load(f)
    except Exception as e:
        raise e

    print(colored("==> Looking for suitable agents in the library...", "green"), flush=True)
    skills = building_task.replace(":", " ").split("\n")
    # skills = [line.split("-", 1)[1].strip() if line.startswith("-") else line for line in lines]
    if len(skills) == 0:
        skills = [building_task]

    chroma_client = chromadb.Client()
    collection = chroma_client.create_collection(
        name="agent_list",
        embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(model_name=embedding_model),
    )
    collection.add(
        documents=[agent["description"] for agent in agent_library],
        metadatas=[{"source": "agent_profile"} for _ in range(len(agent_library))],
        ids=[f"agent_{i}" for i in range(len(agent_library))],
    )
    agent_desc_list = set()
    for skill in skills:
        recall = set(collection.query(query_texts=[skill], n_results=top_k)["documents"][0])
        agent_desc_list = agent_desc_list.union(recall)

    agent_config_list = []
    for description in list(agent_desc_list):
        for agent in agent_library:
            if agent["description"] == description:
                agent_config_list.append(agent.copy())
                break
    chroma_client.delete_collection(collection.name)

    # double recall from the searching result
    expert_pool = [f"{agent['name']}: {agent['description']}" for agent in agent_config_list]
    while True:
        skill_agent_pair_json = (
            self.builder_model.create(
                messages=[
                    {
                        "role": "user",
                        "content": self.AGENT_SELECTION_PROMPT.format(
                            skills=building_task, expert_pool=expert_pool, max_agents=self.max_agents
                        ),
                    }
                ]
            )
            .choices[0]
            .message.content
        )
        try:
            skill_agent_pair_json = _retrieve_json(skill_agent_pair_json)
            skill_agent_pair = json.loads(skill_agent_pair_json)
            break
        except Exception as e:
            print(e, flush=True)
            time.sleep(5)
            continue

    recalled_agent_config_list = []
    recalled_name_desc = []
    for skill, agent_profile in skill_agent_pair.items():
        # If no suitable agent, generate an agent
        if agent_profile == "None":
            _, agent_config_temp = self.build(
                building_task=skill,
                default_llm_config=default_llm_config.copy(),
                coding=False,
                use_oai_assistant=use_oai_assistant,
                max_agents=1,
            )
            self.clear_agent(agent_config_temp["agent_configs"][0]["name"])
            recalled_agent_config_list.append(agent_config_temp["agent_configs"][0])
        else:
            if agent_profile in recalled_name_desc:
                # prevent identical agents
                continue
            recalled_name_desc.append(agent_profile)
            name = agent_profile.split(":")[0].strip()
            desc = agent_profile.split(":")[1].strip()
            for agent in agent_config_list:
                if name == agent["name"] and desc == agent["description"]:
                    recalled_agent_config_list.append(agent.copy())

    print(f"{[agent['name'] for agent in recalled_agent_config_list]} are selected.", flush=True)

    if coding is None:
        resp = (
            self.builder_model.create(
                messages=[{"role": "user", "content": self.CODING_PROMPT.format(task=building_task)}]
            )
            .choices[0]
            .message.content
        )
        coding = resp == "YES"

    self.cached_configs.update({
        "building_task": building_task,
        "agent_configs": recalled_agent_config_list,
        "coding": coding,
        "default_llm_config": default_llm_config,
        "code_execution_config": code_execution_config,
    })
    _config_check(self.cached_configs)

    return self._build_agents(use_oai_assistant, user_proxy=user_proxy, **kwargs)

def save(self, filepath: Optional[str] = None) -> str:
    """Save building configs. If the filepath is not specific, this function will create a filename by encrypt the
    building_task string by md5 with "save_config_" prefix, and save config to the local path.

    Args:
        filepath: save path.

    Return:
        filepath: path save.
    """
    if filepath is None:
        filepath = f"./save_config_{hashlib.md5(self.building_task.encode('utf-8')).hexdigest()}.json"
    with open(filepath, "w") as save_file:
        json.dump(self.cached_configs, save_file, indent=4)
    print(colored(f"Building config saved to {filepath}", "green"), flush=True)

    return filepath

def load(
    self,
    filepath: Optional[str] = None,
    config_json: Optional[str] = None,
    use_oai_assistant: Optional[bool] = False,
    **kwargs: Any,
) -> tuple[list[ConversableAgent], dict[str, Any]]:
    """Load building configs and call the build function to complete building without calling online LLMs' api.

    Args:
        filepath: filepath or JSON string for the save config.
        config_json: JSON string for the save config.
        use_oai_assistant: use OpenAI assistant api instead of self-constructed agent.
        **kwargs (Any): Additional arguments to pass to _build_agents:
            - code_execution_config (Optional[dict[str, Any]]): If provided, overrides the
            code execution configuration from the loaded config.

    Returns:
        agent_list: a list of agents.
        cached_configs: cached configs.
    """
    # load json string.
    if config_json is not None:
        print(colored("Loading config from JSON...", "green"), flush=True)
        cached_configs = json.loads(config_json)

    # load from path.
    if filepath is not None:
        print(colored(f"Loading config from {filepath}", "green"), flush=True)
        with open(filepath) as f:
            cached_configs = json.load(f)

    _config_check(cached_configs)

    agent_configs = cached_configs["agent_configs"]
    default_llm_config = cached_configs["default_llm_config"]
    coding = cached_configs["coding"]

    if kwargs.get("code_execution_config") is not None:
        # for test
        self.cached_configs.update({
            "building_task": cached_configs["building_task"],
            "agent_configs": agent_configs,
            "coding": coding,
            "default_llm_config": default_llm_config,
            "code_execution_config": kwargs["code_execution_config"],
        })
        del kwargs["code_execution_config"]
        return self._build_agents(use_oai_assistant, **kwargs)
    else:
        code_execution_config = cached_configs["code_execution_config"]
        self.cached_configs.update({
            "building_task": cached_configs["building_task"],
            "agent_configs": agent_configs,
            "coding": coding,
            "default_llm_config": default_llm_config,
            "code_execution_config": code_execution_config,
        })
        return self._build_agents(use_oai_assistant, **kwargs)