ForcePilot/src/agents/__init__.py

import asyncio
import importlib
import inspect
from pathlib import Path

from src.agents.common.base import BaseAgent
from src.utils import logger


class AgentManager:
    def __init__(self):
        self._classes = {}
        self._instances = {}  # 存储已创建的 agent 实例

    def register_agent(self, agent_class):
        self._classes[agent_class.__name__] = agent_class

    def init_all_agents(self):
        for agent_id in self._classes.keys():
            self.get_agent(agent_id)

    def get_agent(self, agent_id, reload=False, **kwargs):
        # 检查是否已经创建了该 agent 的实例
        if reload or agent_id not in self._instances:
            agent_class = self._classes[agent_id]
            self._instances[agent_id] = agent_class()

        return self._instances[agent_id]

    def get_agents(self):
        return list(self._instances.values())

    async def reload_all(self):
        for agent_id in self._classes.keys():
            self.get_agent(agent_id, reload=True)

    async def get_agents_info(self):
        agents = self.get_agents()
        return await asyncio.gather(*[a.get_info() for a in agents])

    def auto_discover_agents(self):
        """自动发现并注册 src/agents/ 下的所有智能体。

        遍历 src/agents/ 目录下的所有子文件夹，如果子文件夹包含 __init__.py，
        则尝试从中导入 BaseAgent 的子类并注册。(使用自动导入的方式，支持私有agent)
        """
        # 获取 agents 目录的路径
        agents_dir = Path(__file__).parent

        # 遍历所有子目录
        for item in agents_dir.iterdir():
            logger.info(f"尝试导入模块：{item}")
            # 跳过非目录、common 目录、__pycache__ 等
            if not item.is_dir() or item.name.startswith("_") or item.name == "common":
                continue

            # 检查是否有 __init__.py 文件
            init_file = item / "__init__.py"
            if not init_file.exists():
                logger.warning(f"{item} 不是一个有效的模块")
                continue

            # 尝试导入模块
            try:
                module_name = f"src.agents.{item.name}"
                module = importlib.import_module(module_name)

                # 查找模块中所有 BaseAgent 的子类
                for name, obj in inspect.getmembers(module):
                    if (
                        inspect.isclass(obj)
                        and issubclass(obj, BaseAgent)
                        and obj is not BaseAgent
                        and obj.__module__.startswith(module_name)
                    ):
                        logger.info(f"自动发现智能体: {obj.__name__} 来自 {item.name}")
                        self.register_agent(obj)

            except Exception as e:
                logger.warning(f"无法从 {item.name} 加载智能体: {e}")


agent_manager = AgentManager()
# 自动发现并注册所有智能体
agent_manager.auto_discover_agents()
agent_manager.init_all_agents()

__all__ = ["agent_manager"]


if __name__ == "__main__":
    pass