513 lines
18 KiB
Python
513 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Father Agent - Capability Registry
|
|
|
|
Central registry for tracking capabilities of all child agents (APs).
|
|
Provides capability discovery, matching, and coordination services.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import hashlib
|
|
from dataclasses import dataclass, field, asdict
|
|
from typing import Dict, List, Optional, Any, Callable, Set, Union
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
import threading
|
|
|
|
|
|
class AgentStatus(Enum):
|
|
"""Status of a registered agent."""
|
|
ONLINE = "online"
|
|
OFFLINE = "offline"
|
|
BUSY = "busy"
|
|
DEGRADED = "degraded"
|
|
UNKNOWN = "unknown"
|
|
|
|
|
|
class CapabilityType(Enum):
|
|
"""Types of capabilities."""
|
|
TOOL = "tool"
|
|
SKILL = "skill"
|
|
KNOWLEDGE = "knowledge"
|
|
API = "api"
|
|
PROTOCOL = "protocol"
|
|
SERVICE = "service"
|
|
|
|
|
|
@dataclass
|
|
class CapabilityInfo:
|
|
"""Information about a single capability."""
|
|
name: str
|
|
type: CapabilityType
|
|
version: str
|
|
description: str
|
|
agent_id: str
|
|
parameters: Dict[str, Any] = field(default_factory=dict)
|
|
returns: Dict[str, Any] = field(default_factory=dict)
|
|
tags: List[str] = field(default_factory=list)
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
def __post_init__(self):
|
|
if isinstance(self.type, str):
|
|
self.type = CapabilityType(self.type)
|
|
|
|
|
|
@dataclass
|
|
class AgentCapabilities:
|
|
"""Capabilities registered by an agent."""
|
|
agent_id: str
|
|
endpoint: str
|
|
manifest: Dict[str, Any] = field(default_factory=dict)
|
|
capabilities: List[CapabilityInfo] = field(default_factory=list)
|
|
status: AgentStatus = AgentStatus.UNKNOWN
|
|
registered_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())
|
|
last_seen: str = field(default_factory=lambda: datetime.utcnow().isoformat())
|
|
capability_hash: Optional[str] = None
|
|
|
|
def update_last_seen(self):
|
|
"""Update last seen timestamp."""
|
|
self.last_seen = datetime.utcnow().isoformat()
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
"agent_id": self.agent_id,
|
|
"endpoint": self.endpoint,
|
|
"manifest": self.manifest,
|
|
"status": self.status.value if isinstance(self.status, Enum) else self.status,
|
|
"registered_at": self.registered_at,
|
|
"last_seen": self.last_seen,
|
|
"capability_count": len(self.capabilities),
|
|
"capability_hash": self.capability_hash
|
|
}
|
|
|
|
|
|
class FatherCapabilityRegistry:
|
|
"""
|
|
Central capability registry for father agent.
|
|
|
|
Manages:
|
|
- Agent registration and tracking
|
|
- Capability aggregation
|
|
- Task-capability matching
|
|
- Health monitoring
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._agents: Dict[str, AgentCapabilities] = {}
|
|
self._capabilities_by_name: Dict[str, List[str]] = {} # name -> [agent_ids]
|
|
self._capabilities_by_tag: Dict[str, List[str]] = {} # tag -> [agent_ids]
|
|
self._capabilities_by_type: Dict[CapabilityType, List[str]] = {} # type -> [agent_ids]
|
|
|
|
self._lock = threading.RLock()
|
|
self._listeners: List[Callable[[str, str, Any], None]] = [] # event listeners
|
|
|
|
# Statistics
|
|
self._stats = {
|
|
"registrations": 0,
|
|
"deregistrations": 0,
|
|
"updates": 0,
|
|
"lookups": 0
|
|
}
|
|
|
|
def register_agent(self, agent_info: AgentCapabilities) -> bool:
|
|
"""
|
|
Register a new agent or update existing.
|
|
|
|
Args:
|
|
agent_info: Agent capability information
|
|
|
|
Returns:
|
|
True if new registration, False if update
|
|
"""
|
|
with self._lock:
|
|
agent_id = agent_info.agent_id
|
|
is_new = agent_id not in self._agents
|
|
|
|
# Update indices
|
|
if not is_new:
|
|
self._remove_from_indices(agent_id)
|
|
|
|
self._agents[agent_id] = agent_info
|
|
self._add_to_indices(agent_info)
|
|
|
|
# Update stats
|
|
if is_new:
|
|
self._stats["registrations"] += 1
|
|
self._notify("agent_registered", agent_id, agent_info.to_dict())
|
|
else:
|
|
self._stats["updates"] += 1
|
|
self._notify("agent_updated", agent_id, agent_info.to_dict())
|
|
|
|
return is_new
|
|
|
|
def unregister_agent(self, agent_id: str) -> bool:
|
|
"""Unregister an agent."""
|
|
with self._lock:
|
|
if agent_id not in self._agents:
|
|
return False
|
|
|
|
self._remove_from_indices(agent_id)
|
|
agent_info = self._agents.pop(agent_id)
|
|
|
|
self._stats["deregistrations"] += 1
|
|
self._notify("agent_unregistered", agent_id, agent_info.to_dict())
|
|
|
|
return True
|
|
|
|
def update_agent_status(self, agent_id: str, status: AgentStatus) -> bool:
|
|
"""Update agent status."""
|
|
with self._lock:
|
|
if agent_id not in self._agents:
|
|
return False
|
|
|
|
old_status = self._agents[agent_id].status
|
|
self._agents[agent_id].status = status
|
|
self._agents[agent_id].update_last_seen()
|
|
|
|
self._notify("status_changed", agent_id, {
|
|
"old": old_status.value if isinstance(old_status, Enum) else old_status,
|
|
"new": status.value if isinstance(status, Enum) else status
|
|
})
|
|
|
|
return True
|
|
|
|
def get_agent(self, agent_id: str) -> Optional[AgentCapabilities]:
|
|
"""Get agent capabilities."""
|
|
with self._lock:
|
|
return self._agents.get(agent_id)
|
|
|
|
def list_agents(
|
|
self,
|
|
status: Optional[AgentStatus] = None,
|
|
capability: Optional[str] = None,
|
|
tag: Optional[str] = None
|
|
) -> List[AgentCapabilities]:
|
|
"""
|
|
List agents with optional filters.
|
|
|
|
Args:
|
|
status: Filter by status
|
|
capability: Filter by capability name
|
|
tag: Filter by tag
|
|
|
|
Returns:
|
|
List of matching agents
|
|
"""
|
|
with self._lock:
|
|
self._stats["lookups"] += 1
|
|
|
|
# Start with capability filter if specified
|
|
if capability and capability in self._capabilities_by_name:
|
|
agent_ids = set(self._capabilities_by_name[capability])
|
|
elif tag and tag in self._capabilities_by_tag:
|
|
agent_ids = set(self._capabilities_by_tag[tag])
|
|
else:
|
|
agent_ids = set(self._agents.keys())
|
|
|
|
results = []
|
|
for agent_id in agent_ids:
|
|
agent = self._agents.get(agent_id)
|
|
if not agent:
|
|
continue
|
|
|
|
if status and agent.status != status:
|
|
continue
|
|
|
|
results.append(agent)
|
|
|
|
return results
|
|
|
|
def find_capable_agents(
|
|
self,
|
|
required_capabilities: List[str],
|
|
require_all: bool = True
|
|
) -> List[AgentCapabilities]:
|
|
"""
|
|
Find agents with specified capabilities.
|
|
|
|
Args:
|
|
required_capabilities: List of required capability names
|
|
require_all: If True, agent must have all capabilities
|
|
|
|
Returns:
|
|
List of capable agents
|
|
"""
|
|
with self._lock:
|
|
if not required_capabilities:
|
|
return list(self._agents.values())
|
|
|
|
# Get candidate agent sets
|
|
candidate_sets = []
|
|
for cap in required_capabilities:
|
|
agents_with_cap = set(self._capabilities_by_name.get(cap, []))
|
|
candidate_sets.append(agents_with_cap)
|
|
|
|
if not candidate_sets:
|
|
return []
|
|
|
|
# Combine based on require_all
|
|
if require_all:
|
|
candidate_ids = candidate_sets[0].intersection(*candidate_sets[1:])
|
|
else:
|
|
candidate_ids = candidate_sets[0].union(*candidate_sets[1:])
|
|
|
|
# Filter to online agents
|
|
return [
|
|
self._agents[aid] for aid in candidate_ids
|
|
if aid in self._agents and self._agents[aid].status == AgentStatus.ONLINE
|
|
]
|
|
|
|
def get_capability_owners(self, capability_name: str) -> List[str]:
|
|
"""Get list of agents that have a specific capability."""
|
|
with self._lock:
|
|
return self._capabilities_by_name.get(capability_name, [])
|
|
|
|
def get_all_capabilities(self) -> Dict[str, List[Dict[str, Any]]]:
|
|
"""Get all capabilities grouped by agent."""
|
|
with self._lock:
|
|
return {
|
|
agent_id: [self._cap_to_dict(cap) for cap in agent_info.capabilities]
|
|
for agent_id, agent_info in self._agents.items()
|
|
}
|
|
|
|
def get_capability_catalog(self) -> Dict[str, Any]:
|
|
"""Get catalog of all available capabilities."""
|
|
with self._lock:
|
|
catalog = {
|
|
"by_name": {},
|
|
"by_type": {},
|
|
"by_tag": {}
|
|
}
|
|
|
|
for agent_id, agent_info in self._agents.items():
|
|
for cap in agent_info.capabilities:
|
|
# By name
|
|
if cap.name not in catalog["by_name"]:
|
|
catalog["by_name"][cap.name] = {
|
|
"agents": [],
|
|
"versions": set(),
|
|
"description": cap.description
|
|
}
|
|
catalog["by_name"][cap.name]["agents"].append(agent_id)
|
|
catalog["by_name"][cap.name]["versions"].add(cap.version)
|
|
|
|
# By type
|
|
cap_type = cap.type.value if isinstance(cap.type, Enum) else cap.type
|
|
if cap_type not in catalog["by_type"]:
|
|
catalog["by_type"][cap_type] = []
|
|
if cap.name not in catalog["by_type"][cap_type]:
|
|
catalog["by_type"][cap_type].append(cap.name)
|
|
|
|
# By tag
|
|
for tag in cap.tags:
|
|
if tag not in catalog["by_tag"]:
|
|
catalog["by_tag"][tag] = []
|
|
if cap.name not in catalog["by_tag"][tag]:
|
|
catalog["by_tag"][tag].append(cap.name)
|
|
|
|
# Convert sets to lists for JSON serialization
|
|
for name_info in catalog["by_name"].values():
|
|
name_info["versions"] = list(name_info["versions"])
|
|
|
|
return catalog
|
|
|
|
def match_task_to_agent(
|
|
self,
|
|
task_description: str,
|
|
required_capabilities: List[str],
|
|
preferences: Optional[Dict[str, Any]] = None
|
|
) -> Optional[str]:
|
|
"""
|
|
Match a task to the most suitable agent.
|
|
|
|
Args:
|
|
task_description: Description of the task
|
|
required_capabilities: Required capabilities
|
|
preferences: Optional preferences (e.g., {'min_performance': 0.9})
|
|
|
|
Returns:
|
|
Agent ID of best match or None
|
|
"""
|
|
with self._lock:
|
|
candidates = self.find_capable_agents(required_capabilities)
|
|
|
|
if not candidates:
|
|
return None
|
|
|
|
# Simple scoring based on performance metrics if available
|
|
scored = []
|
|
for agent in candidates:
|
|
score = 1.0 # Base score
|
|
|
|
manifest = agent.manifest
|
|
if "performance" in manifest:
|
|
perf = manifest["performance"]
|
|
score += perf.get("accuracy_score", 0) * 0.5
|
|
score += (perf.get("uptime_percent", 100) / 100) * 0.3
|
|
|
|
scored.append((agent.agent_id, score))
|
|
|
|
# Sort by score
|
|
scored.sort(key=lambda x: x[1], reverse=True)
|
|
|
|
return scored[0][0] if scored else None
|
|
|
|
def add_event_listener(self, callback: Callable[[str, str, Any], None]):
|
|
"""Add an event listener."""
|
|
self._listeners.append(callback)
|
|
|
|
def remove_event_listener(self, callback: Callable[[str, str, Any], None]):
|
|
"""Remove an event listener."""
|
|
if callback in self._listeners:
|
|
self._listeners.remove(callback)
|
|
|
|
def _notify(self, event: str, agent_id: str, data: Any):
|
|
"""Notify event listeners."""
|
|
for listener in self._listeners:
|
|
try:
|
|
listener(event, agent_id, data)
|
|
except Exception as e:
|
|
print(f"Event listener error: {e}")
|
|
|
|
def _add_to_indices(self, agent_info: AgentCapabilities):
|
|
"""Add agent to capability indices."""
|
|
agent_id = agent_info.agent_id
|
|
|
|
for cap in agent_info.capabilities:
|
|
# By name
|
|
if cap.name not in self._capabilities_by_name:
|
|
self._capabilities_by_name[cap.name] = []
|
|
if agent_id not in self._capabilities_by_name[cap.name]:
|
|
self._capabilities_by_name[cap.name].append(agent_id)
|
|
|
|
# By type
|
|
cap_type = cap.type if isinstance(cap.type, CapabilityType) else CapabilityType(cap.type)
|
|
if cap_type not in self._capabilities_by_type:
|
|
self._capabilities_by_type[cap_type] = []
|
|
if agent_id not in self._capabilities_by_type[cap_type]:
|
|
self._capabilities_by_type[cap_type].append(agent_id)
|
|
|
|
# By tag
|
|
for tag in cap.tags:
|
|
if tag not in self._capabilities_by_tag:
|
|
self._capabilities_by_tag[tag] = []
|
|
if agent_id not in self._capabilities_by_tag[tag]:
|
|
self._capabilities_by_tag[tag].append(agent_id)
|
|
|
|
def _remove_from_indices(self, agent_id: str):
|
|
"""Remove agent from capability indices."""
|
|
# Remove from name index
|
|
for name_list in self._capabilities_by_name.values():
|
|
if agent_id in name_list:
|
|
name_list.remove(agent_id)
|
|
|
|
# Remove from type index
|
|
for type_list in self._capabilities_by_type.values():
|
|
if agent_id in type_list:
|
|
type_list.remove(agent_id)
|
|
|
|
# Remove from tag index
|
|
for tag_list in self._capabilities_by_tag.values():
|
|
if agent_id in tag_list:
|
|
tag_list.remove(agent_id)
|
|
|
|
def _cap_to_dict(self, cap: CapabilityInfo) -> Dict[str, Any]:
|
|
"""Convert capability to dictionary."""
|
|
return {
|
|
"name": cap.name,
|
|
"type": cap.type.value if isinstance(cap.type, Enum) else cap.type,
|
|
"version": cap.version,
|
|
"description": cap.description,
|
|
"parameters": cap.parameters,
|
|
"returns": cap.returns,
|
|
"tags": cap.tags
|
|
}
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
"""Get registry statistics."""
|
|
with self._lock:
|
|
return {
|
|
**self._stats,
|
|
"registered_agents": len(self._agents),
|
|
"unique_capabilities": len(self._capabilities_by_name),
|
|
"capability_types": len(self._capabilities_by_type),
|
|
"unique_tags": len(self._capabilities_by_tag),
|
|
"agents_by_status": {
|
|
status.value: sum(
|
|
1 for a in self._agents.values() if a.status == status
|
|
)
|
|
for status in AgentStatus
|
|
}
|
|
}
|
|
|
|
def export(self) -> Dict[str, Any]:
|
|
"""Export registry state."""
|
|
with self._lock:
|
|
return {
|
|
"exported_at": datetime.utcnow().isoformat(),
|
|
"stats": self._stats,
|
|
"agents": {
|
|
agent_id: agent_info.to_dict()
|
|
for agent_id, agent_info in self._agents.items()
|
|
},
|
|
"capability_catalog": self.get_capability_catalog()
|
|
}
|
|
|
|
def export_json(self, indent: int = 2) -> str:
|
|
"""Export as JSON."""
|
|
return json.dumps(self.export(), indent=indent)
|
|
|
|
|
|
# Global registry instance
|
|
_registry_instance: Optional[FatherCapabilityRegistry] = None
|
|
|
|
|
|
def get_father_registry() -> FatherCapabilityRegistry:
|
|
"""Get or create the global father registry."""
|
|
global _registry_instance
|
|
if _registry_instance is None:
|
|
_registry_instance = FatherCapabilityRegistry()
|
|
return _registry_instance
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Demo
|
|
registry = FatherCapabilityRegistry()
|
|
|
|
# Register some agents
|
|
agent1 = AgentCapabilities(
|
|
agent_id="ap-agent-1",
|
|
endpoint="http://localhost:9001",
|
|
capabilities=[
|
|
CapabilityInfo("code_generation", CapabilityType.SKILL, "1.0", "Generate code", "ap-agent-1"),
|
|
CapabilityInfo("file_reader", CapabilityType.TOOL, "1.0", "Read files", "ap-agent-1", tags=["filesystem"])
|
|
],
|
|
status=AgentStatus.ONLINE
|
|
)
|
|
|
|
agent2 = AgentCapabilities(
|
|
agent_id="ap-agent-2",
|
|
endpoint="http://localhost:9002",
|
|
capabilities=[
|
|
CapabilityInfo("code_analysis", CapabilityType.SKILL, "2.0", "Analyze code", "ap-agent-2"),
|
|
CapabilityInfo("file_reader", CapabilityType.TOOL, "1.0", "Read files", "ap-agent-2", tags=["filesystem"])
|
|
],
|
|
status=AgentStatus.ONLINE
|
|
)
|
|
|
|
registry.register_agent(agent1)
|
|
registry.register_agent(agent2)
|
|
|
|
print("Registry stats:", registry.get_stats())
|
|
print("\nCapability catalog:", json.dumps(registry.get_capability_catalog(), indent=2))
|
|
|
|
print("\nAgents with code_generation:", registry.get_capability_owners("code_generation"))
|
|
print("Agents with file_reader:", registry.get_capability_owners("file_reader"))
|
|
|
|
print("\nMatch task to agent:", registry.match_task_to_agent(
|
|
"Generate Python code",
|
|
["code_generation"]
|
|
))
|