171 lines
5.5 KiB
Python
171 lines
5.5 KiB
Python
"""Standalone Runner - 自动发现并启动配置驱动的 Agent
|
||
|
||
扫描 configs/agents/ 目录下的 YAML 文件，自动注册和启动 Agent。
|
||
支持命令行启动：python -m agentkit.core.standalone
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
from agentkit.core.config_driven import AgentConfig, ConfigDrivenAgent
|
||
from agentkit.tools.registry import ToolRegistry
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
DEFAULT_CONFIG_DIR = "configs/agents"
|
||
|
||
|
||
class StandaloneRunner:
|
||
"""自动发现并启动配置驱动的 Agent
|
||
|
||
用法::
|
||
|
||
runner = StandaloneRunner(config_dir="configs/agents")
|
||
runner.add_tool(FunctionTool.from_func(my_tool_func))
|
||
await runner.start_all(redis_url="redis://localhost:6379")
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
config_dir: str = DEFAULT_CONFIG_DIR,
|
||
tool_registry: ToolRegistry | None = None,
|
||
llm_client=None,
|
||
custom_handlers: dict | None = None,
|
||
):
|
||
self._config_dir = config_dir
|
||
self._tool_registry = tool_registry or ToolRegistry()
|
||
self._llm_client = llm_client
|
||
self._custom_handlers = custom_handlers or {}
|
||
self._agents: dict[str, ConfigDrivenAgent] = {}
|
||
|
||
@property
|
||
def agents(self) -> dict[str, ConfigDrivenAgent]:
|
||
return self._agents
|
||
|
||
def add_tool(self, tool) -> "StandaloneRunner":
|
||
"""添加工具到注册中心"""
|
||
self._tool_registry.register(tool)
|
||
return self
|
||
|
||
def add_custom_handler(self, name: str, handler) -> "StandaloneRunner":
|
||
"""注册自定义 handler"""
|
||
self._custom_handlers[name] = handler
|
||
return self
|
||
|
||
def discover_configs(self) -> list[AgentConfig]:
|
||
"""扫描配置目录，发现所有 YAML 配置"""
|
||
configs = []
|
||
config_path = Path(self._config_dir)
|
||
|
||
if not config_path.exists():
|
||
logger.warning(f"Config directory '{self._config_dir}' not found")
|
||
return configs
|
||
|
||
for yaml_file in sorted(config_path.glob("*.yaml")):
|
||
try:
|
||
config = AgentConfig.from_yaml(str(yaml_file))
|
||
configs.append(config)
|
||
logger.info(f"Discovered agent config: {config.name} from {yaml_file.name}")
|
||
except Exception as e:
|
||
logger.error(f"Failed to load config '{yaml_file}': {e}")
|
||
|
||
for yml_file in sorted(config_path.glob("*.yml")):
|
||
try:
|
||
config = AgentConfig.from_yaml(str(yml_file))
|
||
configs.append(config)
|
||
logger.info(f"Discovered agent config: {config.name} from {yml_file.name}")
|
||
except Exception as e:
|
||
logger.error(f"Failed to load config '{yml_file}': {e}")
|
||
|
||
return configs
|
||
|
||
def build_agents(self) -> dict[str, ConfigDrivenAgent]:
|
||
"""从发现的配置构建 Agent 实例"""
|
||
configs = self.discover_configs()
|
||
self._agents.clear()
|
||
|
||
for config in configs:
|
||
try:
|
||
agent = ConfigDrivenAgent(
|
||
config=config,
|
||
tool_registry=self._tool_registry,
|
||
llm_client=self._llm_client,
|
||
custom_handlers=self._custom_handlers,
|
||
)
|
||
self._agents[config.name] = agent
|
||
logger.info(f"Built agent: {config.name} (mode={config.task_mode})")
|
||
except Exception as e:
|
||
logger.error(f"Failed to build agent '{config.name}': {e}")
|
||
|
||
return self._agents
|
||
|
||
async def start_all(self, redis_url: str = "") -> None:
|
||
"""启动所有已构建的 Agent"""
|
||
if not self._agents:
|
||
self.build_agents()
|
||
|
||
for name, agent in self._agents.items():
|
||
try:
|
||
await agent.start(redis_url=redis_url)
|
||
logger.info(f"Agent '{name}' started")
|
||
except Exception as e:
|
||
logger.error(f"Failed to start agent '{name}': {e}")
|
||
|
||
async def stop_all(self) -> None:
|
||
"""停止所有 Agent"""
|
||
for name, agent in self._agents.items():
|
||
try:
|
||
await agent.stop()
|
||
logger.info(f"Agent '{name}' stopped")
|
||
except Exception as e:
|
||
logger.error(f"Failed to stop agent '{name}': {e}")
|
||
|
||
async def execute_task(self, agent_name: str, task) -> dict | None:
|
||
"""在指定 Agent 上执行任务（本地模式）"""
|
||
if agent_name not in self._agents:
|
||
logger.error(f"Agent '{agent_name}' not found")
|
||
return None
|
||
|
||
agent = self._agents[agent_name]
|
||
result = await agent.execute(task)
|
||
return result.to_dict()
|
||
|
||
|
||
async def main():
|
||
"""命令行入口"""
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
|
||
)
|
||
|
||
config_dir = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_CONFIG_DIR
|
||
redis_url = os.environ.get("REDIS_URL", "")
|
||
|
||
runner = StandaloneRunner(config_dir=config_dir)
|
||
agents = runner.build_agents()
|
||
|
||
if not agents:
|
||
logger.error("No agents discovered. Check your config directory.")
|
||
sys.exit(1)
|
||
|
||
logger.info(f"Discovered {len(agents)} agent(s): {list(agents.keys())}")
|
||
|
||
try:
|
||
await runner.start_all(redis_url=redis_url)
|
||
logger.info("All agents started. Press Ctrl+C to stop.")
|
||
|
||
# 保持运行
|
||
while True:
|
||
await asyncio.sleep(1)
|
||
except KeyboardInterrupt:
|
||
logger.info("Shutting down...")
|
||
finally:
|
||
await runner.stop_all()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|