251 lines
9.6 KiB
Python
251 lines
9.6 KiB
Python
"""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,
|
||
}
|