478 lines
15 KiB
Python
478 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
AP Capability Registry Server
|
|
|
|
Server for advertising capabilities to father and querying other agents.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass, field, asdict
|
|
from typing import Dict, List, Optional, Any, Callable, Set
|
|
from datetime import datetime, timedelta
|
|
from enum import Enum
|
|
import threading
|
|
import time
|
|
|
|
from discovery import CapabilityRegistry, Capability, CapabilityQuery, CapabilityType, get_registry
|
|
from manifest import AgentManifest, get_default_manifest
|
|
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("ap-capability-server")
|
|
|
|
|
|
class AgentStatus(Enum):
|
|
"""Status of a registered agent."""
|
|
ONLINE = "online"
|
|
OFFLINE = "offline"
|
|
BUSY = "busy"
|
|
UNKNOWN = "unknown"
|
|
|
|
|
|
@dataclass
|
|
class AgentInfo:
|
|
"""Information about a registered agent."""
|
|
agent_id: str
|
|
manifest: AgentManifest
|
|
endpoint: str
|
|
status: AgentStatus = AgentStatus.ONLINE
|
|
last_seen: str = field(default_factory=lambda: datetime.utcnow().isoformat())
|
|
capabilities_hash: Optional[str] = None
|
|
|
|
def update_last_seen(self):
|
|
"""Update last seen timestamp."""
|
|
self.last_seen = datetime.utcnow().isoformat()
|
|
|
|
|
|
class CapabilityServer:
|
|
"""
|
|
Server for managing capability advertisement and discovery.
|
|
|
|
This allows AP to:
|
|
- Advertise its capabilities to father/other agents
|
|
- Query capabilities of other agents
|
|
- Match tasks to capable agents
|
|
"""
|
|
|
|
def __init__(self, agent_id: str, father_endpoint: Optional[str] = None):
|
|
self.agent_id = agent_id
|
|
self.father_endpoint = father_endpoint
|
|
self.registry = get_registry(agent_id)
|
|
self.manifest = get_default_manifest()
|
|
|
|
# Known agents
|
|
self._agents: Dict[str, AgentInfo] = {}
|
|
self._agents_lock = threading.RLock()
|
|
|
|
# Task matching callbacks
|
|
self._match_callbacks: List[Callable[[str, Dict[str, Any]], Optional[str]]] = []
|
|
|
|
# Server state
|
|
self._running = False
|
|
self._heartbeat_task = None
|
|
self._cleanup_task = None
|
|
|
|
# Configuration
|
|
self.heartbeat_interval = 30 # seconds
|
|
self.agent_timeout = 120 # seconds
|
|
|
|
def start(self):
|
|
"""Start the capability server."""
|
|
if self._running:
|
|
return
|
|
|
|
self._running = True
|
|
logger.info(f"Starting capability server for {self.agent_id}")
|
|
|
|
# Start background tasks
|
|
self._heartbeat_task = threading.Thread(target=self._heartbeat_loop, daemon=True)
|
|
self._heartbeat_task.start()
|
|
|
|
self._cleanup_task = threading.Thread(target=self._cleanup_loop, daemon=True)
|
|
self._cleanup_task.start()
|
|
|
|
# Advertise to father if configured
|
|
if self.father_endpoint:
|
|
self.advertise_to_father()
|
|
|
|
def stop(self):
|
|
"""Stop the capability server."""
|
|
logger.info(f"Stopping capability server for {self.agent_id}")
|
|
self._running = False
|
|
|
|
def _heartbeat_loop(self):
|
|
"""Background loop for sending heartbeats."""
|
|
while self._running:
|
|
try:
|
|
self._send_heartbeat()
|
|
time.sleep(self.heartbeat_interval)
|
|
except Exception as e:
|
|
logger.error(f"Heartbeat error: {e}")
|
|
time.sleep(5)
|
|
|
|
def _cleanup_loop(self):
|
|
"""Background loop for cleaning up stale agents."""
|
|
while self._running:
|
|
try:
|
|
self._cleanup_stale_agents()
|
|
time.sleep(self.heartbeat_interval * 2)
|
|
except Exception as e:
|
|
logger.error(f"Cleanup error: {e}")
|
|
time.sleep(5)
|
|
|
|
def _send_heartbeat(self):
|
|
"""Send heartbeat to father and update last seen."""
|
|
# In a real implementation, this would make an HTTP/WebSocket call
|
|
# For now, we just update our local state
|
|
pass
|
|
|
|
def _cleanup_stale_agents(self):
|
|
"""Remove agents that haven't been seen recently."""
|
|
with self._agents_lock:
|
|
now = datetime.utcnow()
|
|
stale = []
|
|
|
|
for agent_id, info in self._agents.items():
|
|
last_seen = datetime.fromisoformat(info.last_seen)
|
|
if (now - last_seen).total_seconds() > self.agent_timeout:
|
|
stale.append(agent_id)
|
|
|
|
for agent_id in stale:
|
|
logger.info(f"Removing stale agent: {agent_id}")
|
|
del self._agents[agent_id]
|
|
|
|
def register_agent(self, agent_info: AgentInfo) -> bool:
|
|
"""
|
|
Register another agent.
|
|
|
|
Args:
|
|
agent_info: Information about the agent
|
|
|
|
Returns:
|
|
True if registered, False if updated existing
|
|
"""
|
|
with self._agents_lock:
|
|
is_new = agent_info.agent_id not in self._agents
|
|
agent_info.update_last_seen()
|
|
self._agents[agent_info.agent_id] = agent_info
|
|
|
|
if is_new:
|
|
logger.info(f"Registered new agent: {agent_info.agent_id}")
|
|
else:
|
|
logger.debug(f"Updated agent: {agent_info.agent_id}")
|
|
|
|
return is_new
|
|
|
|
def unregister_agent(self, agent_id: str) -> bool:
|
|
"""Unregister an agent."""
|
|
with self._agents_lock:
|
|
if agent_id in self._agents:
|
|
del self._agents[agent_id]
|
|
logger.info(f"Unregistered agent: {agent_id}")
|
|
return True
|
|
return False
|
|
|
|
def get_agent(self, agent_id: str) -> Optional[AgentInfo]:
|
|
"""Get information about a registered agent."""
|
|
with self._agents_lock:
|
|
return self._agents.get(agent_id)
|
|
|
|
def list_agents(self, status: Optional[AgentStatus] = None) -> List[AgentInfo]:
|
|
"""
|
|
List all registered agents.
|
|
|
|
Args:
|
|
status: Optional filter by status
|
|
|
|
Returns:
|
|
List of agent information
|
|
"""
|
|
with self._agents_lock:
|
|
agents = list(self._agents.values())
|
|
if status:
|
|
agents = [a for a in agents if a.status == status]
|
|
return agents
|
|
|
|
def update_agent_status(self, agent_id: str, status: AgentStatus) -> bool:
|
|
"""Update an agent's status."""
|
|
with self._agents_lock:
|
|
if agent_id in self._agents:
|
|
self._agents[agent_id].status = status
|
|
self._agents[agent_id].update_last_seen()
|
|
return True
|
|
return False
|
|
|
|
def advertise_to_father(self) -> bool:
|
|
"""
|
|
Advertise capabilities to father agent.
|
|
|
|
Returns:
|
|
True if advertisement successful
|
|
"""
|
|
if not self.father_endpoint:
|
|
logger.warning("No father endpoint configured")
|
|
return False
|
|
|
|
try:
|
|
# Build advertisement payload
|
|
advertisement = {
|
|
"agent_id": self.agent_id,
|
|
"manifest": self.manifest.to_dict(),
|
|
"capabilities": self.registry.export(),
|
|
"timestamp": datetime.utcnow().isoformat()
|
|
}
|
|
|
|
# In a real implementation, this would be an HTTP POST
|
|
# response = requests.post(
|
|
# f"{self.father_endpoint}/api/agents/register",
|
|
# json=advertisement
|
|
# )
|
|
# return response.status_code == 200
|
|
|
|
logger.info(f"Advertised to father: {self.father_endpoint}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to advertise to father: {e}")
|
|
return False
|
|
|
|
def query_other_agent(self, agent_id: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Query capabilities of another agent.
|
|
|
|
Args:
|
|
agent_id: ID of the agent to query
|
|
|
|
Returns:
|
|
Agent capabilities or None if not found
|
|
"""
|
|
with self._agents_lock:
|
|
agent = self._agents.get(agent_id)
|
|
if not agent:
|
|
return None
|
|
|
|
return {
|
|
"agent_id": agent.agent_id,
|
|
"manifest": agent.manifest.to_dict(),
|
|
"status": agent.status.value,
|
|
"last_seen": agent.last_seen
|
|
}
|
|
|
|
def find_agents_by_capability(
|
|
self,
|
|
capability_name: str,
|
|
min_version: Optional[str] = None
|
|
) -> List[AgentInfo]:
|
|
"""
|
|
Find agents that have a specific capability.
|
|
|
|
Args:
|
|
capability_name: Name of the capability
|
|
min_version: Minimum version required
|
|
|
|
Returns:
|
|
List of capable agents
|
|
"""
|
|
with self._agents_lock:
|
|
capable = []
|
|
|
|
for agent in self._agents.values():
|
|
manifest = agent.manifest
|
|
|
|
# Check tools
|
|
if capability_name in manifest.tools:
|
|
capable.append(agent)
|
|
continue
|
|
|
|
# Check skills
|
|
if capability_name in manifest.skills:
|
|
capable.append(agent)
|
|
continue
|
|
|
|
# Check knowledge areas
|
|
if capability_name in manifest.knowledge_areas:
|
|
capable.append(agent)
|
|
|
|
return capable
|
|
|
|
def find_agents_by_tag(self, tag: str) -> List[AgentInfo]:
|
|
"""Find agents by tag."""
|
|
with self._agents_lock:
|
|
return [
|
|
agent for agent in self._agents.values()
|
|
if tag in agent.manifest.tags
|
|
]
|
|
|
|
def match_task_to_agent(
|
|
self,
|
|
task_type: str,
|
|
requirements: Dict[str, Any]
|
|
) -> Optional[str]:
|
|
"""
|
|
Match a task to the most suitable agent.
|
|
|
|
Args:
|
|
task_type: Type of task
|
|
requirements: Task requirements (capabilities, resources, etc.)
|
|
|
|
Returns:
|
|
Agent ID of best match or None
|
|
"""
|
|
# Check custom matchers first
|
|
for callback in self._match_callbacks:
|
|
result = callback(task_type, requirements)
|
|
if result:
|
|
return result
|
|
|
|
# Default matching logic
|
|
with self._agents_lock:
|
|
candidates = []
|
|
|
|
for agent in self._agents.values():
|
|
if agent.status != AgentStatus.ONLINE:
|
|
continue
|
|
|
|
# Check capability requirements
|
|
required_caps = requirements.get("capabilities", [])
|
|
manifest = agent.manifest
|
|
|
|
has_all_caps = all(
|
|
cap in manifest.tools or
|
|
cap in manifest.skills or
|
|
cap in manifest.knowledge_areas
|
|
for cap in required_caps
|
|
)
|
|
|
|
if has_all_caps:
|
|
# Score based on performance
|
|
score = manifest.performance.accuracy_score
|
|
candidates.append((agent.agent_id, score))
|
|
|
|
if not candidates:
|
|
return None
|
|
|
|
# Return highest scoring agent
|
|
candidates.sort(key=lambda x: x[1], reverse=True)
|
|
return candidates[0][0]
|
|
|
|
def register_match_callback(
|
|
self,
|
|
callback: Callable[[str, Dict[str, Any]], Optional[str]]
|
|
):
|
|
"""
|
|
Register a custom task matching callback.
|
|
|
|
The callback receives (task_type, requirements) and should
|
|
return an agent_id or None.
|
|
"""
|
|
self._match_callbacks.append(callback)
|
|
|
|
def get_system_status(self) -> Dict[str, Any]:
|
|
"""Get overall system status."""
|
|
with self._agents_lock:
|
|
return {
|
|
"server": {
|
|
"agent_id": self.agent_id,
|
|
"running": self._running,
|
|
"father_endpoint": self.father_endpoint
|
|
},
|
|
"agents": {
|
|
"total": len(self._agents),
|
|
"by_status": {
|
|
status.value: sum(
|
|
1 for a in self._agents.values()
|
|
if a.status == status
|
|
)
|
|
for status in AgentStatus
|
|
},
|
|
"list": [
|
|
{
|
|
"agent_id": a.agent_id,
|
|
"status": a.status.value,
|
|
"last_seen": a.last_seen
|
|
}
|
|
for a in self._agents.values()
|
|
]
|
|
},
|
|
"local_capabilities": self.registry.get_stats()
|
|
}
|
|
|
|
def export_capabilities(self, format: str = "json") -> str:
|
|
"""Export all local capabilities."""
|
|
if format == "yaml":
|
|
try:
|
|
import yaml
|
|
return yaml.dump(self.registry.export(), default_flow_style=False)
|
|
except ImportError:
|
|
pass
|
|
return self.registry.export_json()
|
|
|
|
|
|
# Global server instance
|
|
_server_instance: Optional[CapabilityServer] = None
|
|
|
|
|
|
def get_server(
|
|
agent_id: Optional[str] = None,
|
|
father_endpoint: Optional[str] = None
|
|
) -> CapabilityServer:
|
|
"""Get or create the global server instance."""
|
|
global _server_instance
|
|
if _server_instance is None:
|
|
if agent_id is None:
|
|
agent_id = "ap-agent"
|
|
_server_instance = CapabilityServer(agent_id, father_endpoint)
|
|
return _server_instance
|
|
|
|
|
|
def reset_server():
|
|
"""Reset the global server instance."""
|
|
global _server_instance
|
|
if _server_instance:
|
|
_server_instance.stop()
|
|
_server_instance = None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Demo
|
|
server = CapabilityServer("test-agent", "http://localhost:8080")
|
|
|
|
# Register some capabilities
|
|
from discovery import Capability, CapabilityType
|
|
|
|
server.registry.register(Capability(
|
|
name="code_analysis",
|
|
type=CapabilityType.SKILL,
|
|
version="1.0.0",
|
|
description="Analyze code for bugs and improvements"
|
|
))
|
|
|
|
server.start()
|
|
|
|
# Register a mock agent
|
|
mock_agent = AgentInfo(
|
|
agent_id="other-agent",
|
|
manifest=get_default_manifest(),
|
|
endpoint="http://localhost:9000"
|
|
)
|
|
server.register_agent(mock_agent)
|
|
|
|
print("System status:", json.dumps(server.get_system_status(), indent=2))
|
|
|
|
# Find agents by capability
|
|
capable = server.find_agents_by_capability("code_generation")
|
|
print(f"\nAgents with 'code_generation': {[a.agent_id for a in capable]}")
|
|
|
|
# Match task
|
|
best_agent = server.match_task_to_agent("coding", {
|
|
"capabilities": ["code_generation"]
|
|
})
|
|
print(f"\nBest agent for coding task: {best_agent}")
|
|
|
|
server.stop()
|