174 lines
6.1 KiB
Python
174 lines
6.1 KiB
Python
"""OrganizationContext - 组织上下文，管理 AgentProfile 与能力矩阵"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from dataclasses import dataclass, field
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class AgentProfile:
|
||
"""Agent 档案 - 描述组织中一个 Agent 的能力与状态"""
|
||
|
||
name: str
|
||
agent_type: str # "react", "rewoo", "plan_exec", "reflexion", "direct"
|
||
capabilities: list[str] # capability tag strings
|
||
skills: list[str] # skill names
|
||
current_load: int = 0 # number of active tasks
|
||
max_concurrency: int = 1
|
||
availability: bool = True
|
||
specializations: list[str] = field(default_factory=list)
|
||
model: str = "default"
|
||
execution_mode: str = "react"
|
||
|
||
|
||
class OrganizationContext:
|
||
"""组织上下文 - 管理 Agent 档案与能力矩阵，支持基于能力的 Agent 发现"""
|
||
|
||
def __init__(self) -> None:
|
||
self._agents: dict[str, AgentProfile] = {}
|
||
self._capability_matrix: dict[str, list[str]] = {} # capability -> [agent_names]
|
||
|
||
def register_agent(self, profile: AgentProfile) -> None:
|
||
"""注册 Agent 档案"""
|
||
self._agents[profile.name] = profile
|
||
# 更新能力矩阵
|
||
for cap in profile.capabilities:
|
||
cap_lower = cap.lower()
|
||
if cap_lower not in self._capability_matrix:
|
||
self._capability_matrix[cap_lower] = []
|
||
if profile.name not in self._capability_matrix[cap_lower]:
|
||
self._capability_matrix[cap_lower].append(profile.name)
|
||
logger.info(f"Agent profile '{profile.name}' registered")
|
||
|
||
def unregister_agent(self, name: str) -> None:
|
||
"""注销 Agent 档案"""
|
||
profile = self._agents.pop(name, None)
|
||
if profile is None:
|
||
return
|
||
# 清理能力矩阵
|
||
for cap in profile.capabilities:
|
||
cap_lower = cap.lower()
|
||
if cap_lower in self._capability_matrix:
|
||
self._capability_matrix[cap_lower] = [
|
||
n for n in self._capability_matrix[cap_lower] if n != name
|
||
]
|
||
if not self._capability_matrix[cap_lower]:
|
||
del self._capability_matrix[cap_lower]
|
||
logger.info(f"Agent profile '{name}' unregistered")
|
||
|
||
def get_agent_profile(self, name: str) -> AgentProfile | None:
|
||
"""获取 Agent 档案"""
|
||
return self._agents.get(name)
|
||
|
||
def list_agents(self) -> list[AgentProfile]:
|
||
"""列出所有 Agent 档案"""
|
||
return list(self._agents.values())
|
||
|
||
def find_best_agent(
|
||
self,
|
||
required_capabilities: list[str],
|
||
exclude: list[str] | None = None,
|
||
) -> AgentProfile | None:
|
||
"""根据能力需求找到最佳 Agent
|
||
|
||
逻辑：
|
||
1. 找到拥有所有所需能力的 Agent
|
||
2. 在匹配的 Agent 中，优先选择 current_load 较低的
|
||
3. 排除 exclude 列表中的 Agent
|
||
4. 排除不可用的 Agent
|
||
5. 没有匹配则返回 None
|
||
"""
|
||
exclude_set = set(exclude or [])
|
||
|
||
# 对每个所需能力，查找拥有该能力的 Agent 名称集合
|
||
candidate_names: set[str] | None = None
|
||
for cap in required_capabilities:
|
||
cap_lower = cap.lower()
|
||
agents_with_cap = set(self._capability_matrix.get(cap_lower, []))
|
||
if candidate_names is None:
|
||
candidate_names = agents_with_cap
|
||
else:
|
||
candidate_names &= agents_with_cap
|
||
|
||
if not candidate_names:
|
||
return None
|
||
|
||
# 过滤排除和不可用的 Agent，按 load 排序
|
||
candidates = [
|
||
self._agents[name]
|
||
for name in candidate_names
|
||
if name not in exclude_set
|
||
and name in self._agents
|
||
and self._agents[name].availability
|
||
]
|
||
|
||
if not candidates:
|
||
return None
|
||
|
||
candidates.sort(key=lambda p: p.current_load)
|
||
return candidates[0]
|
||
|
||
def update_load(self, name: str, delta: int) -> None:
|
||
"""更新 Agent 负载"""
|
||
profile = self._agents.get(name)
|
||
if profile is not None:
|
||
profile.current_load = max(0, profile.current_load + delta)
|
||
|
||
def set_availability(self, name: str, available: bool) -> None:
|
||
"""设置 Agent 可用性"""
|
||
profile = self._agents.get(name)
|
||
if profile is not None:
|
||
profile.availability = available
|
||
|
||
@classmethod
|
||
def from_agent_pool(cls, agent_pool, skill_registry) -> OrganizationContext:
|
||
"""从 AgentPool 和 SkillRegistry 构建 OrganizationContext
|
||
|
||
Args:
|
||
agent_pool: AgentPool 实例，提供运行时 Agent 列表
|
||
skill_registry: SkillRegistry 实例，提供 Skill 配置查询
|
||
"""
|
||
ctx = cls()
|
||
|
||
if agent_pool is None or skill_registry is None:
|
||
return ctx
|
||
|
||
for agent_info in agent_pool.list_agents():
|
||
agent_name = agent_info["name"]
|
||
agent_type = agent_info.get("agent_type", "react")
|
||
|
||
# 尝试从 skill_registry 获取 SkillConfig
|
||
capabilities: list[str] = []
|
||
skills: list[str] = []
|
||
execution_mode = "react"
|
||
model = "default"
|
||
max_concurrency = 1
|
||
|
||
try:
|
||
skill = skill_registry.get(agent_name)
|
||
config = skill.config
|
||
capabilities = [cap.tag for cap in config.capabilities]
|
||
execution_mode = config.execution_mode
|
||
model = config.llm.get("model", "default") if config.llm else "default"
|
||
max_concurrency = config.max_concurrency
|
||
skills = [agent_name]
|
||
except Exception:
|
||
# Agent 不在 skill_registry 中，使用默认值
|
||
skills = [agent_name]
|
||
|
||
profile = AgentProfile(
|
||
name=agent_name,
|
||
agent_type=agent_type,
|
||
capabilities=capabilities,
|
||
skills=skills,
|
||
execution_mode=execution_mode,
|
||
model=model,
|
||
max_concurrency=max_concurrency,
|
||
)
|
||
ctx.register_agent(profile)
|
||
|
||
return ctx
|