"""Agent 注册中心 - 管理 Agent 的注册、发现、状态 与业务系统解耦：通过 session_factory 注入数据库会话， 通过 agent_model_factory 注入 ORM 模型。 """ import logging from datetime import datetime, timedelta, timezone from typing import Any, Callable, Awaitable from agentkit.core.exceptions import ( AgentNotFoundError, AgentUnavailableError, NoAvailableAgentError, ) from agentkit.core.protocol import AgentCapability, AgentStatus logger = logging.getLogger(__name__) HEARTBEAT_TIMEOUT_SECONDS = 90 class AgentRegistry: """Agent 注册中心，管理 Agent 的注册、发现、状态 使用依赖注入模式，不依赖具体的 ORM 模型或数据库连接。 """ def __init__( self, session_factory: Callable[[], Any], agent_model: Any, load_balancer: str = "round_robin", ): """ Args: session_factory: 返回 async context manager 的工厂，用于获取数据库会话 agent_model: Agent ORM 模型类 load_balancer: 负载均衡策略 (round_robin / least_tasks / random) """ self._session_factory = session_factory self._agent_model = agent_model self._load_balancer = load_balancer self._round_robin_index: dict[str, int] = {} async def register(self, capability: AgentCapability, endpoint: str) -> str: """注册 Agent，返回 agent_id。同名 Agent 已存在则更新。""" async with self._session_factory() as db: try: Model = self._agent_model stmt = type(db).execute.__self__.__class__ # placeholder # 尝试查找已有记录 from sqlalchemy import select stmt = select(Model).where(Model.name == capability.agent_name) result = await db.execute(stmt) existing = result.scalar_one_or_none() if existing: existing.agent_type = capability.agent_type existing.version = capability.version existing.endpoint = endpoint existing.description = capability.description existing.capabilities = capability.to_dict() existing.status = AgentStatus.ONLINE existing.last_heartbeat = datetime.now(timezone.utc) await db.commit() await db.refresh(existing) agent_id = existing.id logger.info(f"Agent '{capability.agent_name}' re-registered (id={agent_id})") else: agent = Model( name=capability.agent_name, display_name=capability.agent_name.replace("_", " ").title(), agent_type=capability.agent_type, description=capability.description, version=capability.version, endpoint=endpoint, status=AgentStatus.ONLINE, capabilities=capability.to_dict(), last_heartbeat=datetime.now(timezone.utc), ) db.add(agent) await db.commit() await db.refresh(agent) agent_id = agent.id logger.info(f"Agent '{capability.agent_name}' registered (id={agent_id})") return str(agent_id) except Exception as e: await db.rollback() logger.error(f"Failed to register agent '{capability.agent_name}': {e}") raise async def unregister(self, agent_name: str): """注销 Agent（设置状态为 offline）""" async with self._session_factory() as db: try: from sqlalchemy import select Model = self._agent_model stmt = select(Model).where(Model.name == agent_name) result = await db.execute(stmt) agent = result.scalar_one_or_none() if not agent: logger.warning(f"Attempted to unregister non-existent agent '{agent_name}'") return agent.status = AgentStatus.OFFLINE await db.commit() logger.info(f"Agent '{agent_name}' unregistered") except Exception as e: await db.rollback() logger.error(f"Failed to unregister agent '{agent_name}': {e}") raise async def update_heartbeat(self, agent_name: str): """更新心跳时间""" async with self._session_factory() as db: try: from sqlalchemy import update Model = self._agent_model stmt = ( update(Model) .where(Model.name == agent_name) .values( last_heartbeat=datetime.now(timezone.utc), status=AgentStatus.ONLINE, ) ) await db.execute(stmt) await db.commit() except Exception as e: await db.rollback() logger.error(f"Failed to update heartbeat for agent '{agent_name}': {e}") async def get_agent(self, agent_name: str) -> dict | None: """获取 Agent 信息""" async with self._session_factory() as db: from sqlalchemy import select Model = self._agent_model stmt = select(Model).where(Model.name == agent_name) result = await db.execute(stmt) agent = result.scalar_one_or_none() if not agent: return None return self._agent_to_dict(agent) async def list_agents( self, agent_type: str | None = None, status: str | None = None, ) -> list[dict]: """列出 Agent，支持按类型和状态筛选""" async with self._session_factory() as db: from sqlalchemy import select Model = self._agent_model stmt = select(Model) if agent_type: stmt = stmt.where(Model.agent_type == agent_type) if status: stmt = stmt.where(Model.status == status) stmt = stmt.order_by(Model.created_at.desc()) result = await db.execute(stmt) agents = result.scalars().all() return [self._agent_to_dict(a) for a in agents] async def get_available_agent(self, task_type: str) -> str | None: """根据任务类型找到可用 Agent（支持负载均衡）""" async with self._session_factory() as db: from sqlalchemy import select Model = self._agent_model stmt = select(Model).where(Model.status == AgentStatus.ONLINE) result = await db.execute(stmt) agents = result.scalars().all() candidates = [] for agent in agents: capabilities = agent.capabilities or {} supported_tasks = capabilities.get("supported_tasks", []) if task_type in supported_tasks: candidates.append(agent) if not candidates: return None # 负载均衡选择 if self._load_balancer == "round_robin": idx = self._round_robin_index.get(task_type, 0) selected = candidates[idx % len(candidates)] self._round_robin_index[task_type] = idx + 1 return selected.name elif self._load_balancer == "random": import random return random.choice(candidates).name else: # least_tasks 或默认：返回第一个 return candidates[0].name async def check_health(self): """检查所有 Agent 健康状态，超时标记为 offline""" async with self._session_factory() as db: try: from sqlalchemy import update Model = self._agent_model timeout_threshold = datetime.now(timezone.utc) - timedelta( seconds=HEARTBEAT_TIMEOUT_SECONDS ) stmt = ( update(Model) .where( Model.status == AgentStatus.ONLINE, Model.last_heartbeat < timeout_threshold, ) .values(status=AgentStatus.OFFLINE) ) result = await db.execute(stmt) await db.commit() if result.rowcount > 0: logger.warning( f"Marked {result.rowcount} agent(s) as offline due to heartbeat timeout" ) except Exception as e: await db.rollback() logger.error(f"Failed to check agent health: {e}") def _agent_to_dict(self, agent: Any) -> dict: """将 Agent ORM 对象转换为字典""" return { "id": str(agent.id), "name": agent.name, "display_name": agent.display_name, "agent_type": agent.agent_type, "description": agent.description, "version": agent.version, "endpoint": agent.endpoint, "status": agent.status, "capabilities": agent.capabilities, "last_heartbeat": agent.last_heartbeat.isoformat() if agent.last_heartbeat else None, "created_at": agent.created_at.isoformat() if agent.created_at else None, "updated_at": agent.updated_at.isoformat() if agent.updated_at else None, }