383 lines
12 KiB
Python
383 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
AP Capability Registry
|
|
|
|
Manages registration, discovery, and versioning of AP capabilities.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import hashlib
|
|
from dataclasses import dataclass, field, asdict
|
|
from typing import Dict, List, Optional, Set, Any, Callable
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
import threading
|
|
|
|
|
|
class CapabilityType(Enum):
|
|
"""Types of capabilities an agent can have."""
|
|
TOOL = "tool"
|
|
SKILL = "skill"
|
|
KNOWLEDGE = "knowledge"
|
|
API = "api"
|
|
PROTOCOL = "protocol"
|
|
SERVICE = "service"
|
|
|
|
|
|
class CapabilityStatus(Enum):
|
|
"""Status of a capability."""
|
|
ACTIVE = "active"
|
|
DEPRECATED = "deprecated"
|
|
EXPERIMENTAL = "experimental"
|
|
DISABLED = "disabled"
|
|
|
|
|
|
@dataclass
|
|
class Capability:
|
|
"""Represents a single capability."""
|
|
name: str
|
|
type: CapabilityType
|
|
version: str
|
|
description: str
|
|
endpoint: Optional[str] = None
|
|
parameters: Dict[str, Any] = field(default_factory=dict)
|
|
returns: Dict[str, Any] = field(default_factory=dict)
|
|
dependencies: List[str] = field(default_factory=list)
|
|
tags: List[str] = field(default_factory=list)
|
|
status: CapabilityStatus = CapabilityStatus.ACTIVE
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
created_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())
|
|
updated_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())
|
|
|
|
def __post_init__(self):
|
|
if isinstance(self.type, str):
|
|
self.type = CapabilityType(self.type)
|
|
if isinstance(self.status, str):
|
|
self.status = CapabilityStatus(self.status)
|
|
|
|
@property
|
|
def id(self) -> str:
|
|
"""Generate unique ID for this capability."""
|
|
content = f"{self.name}:{self.version}:{self.type.value}"
|
|
return hashlib.sha256(content.encode()).hexdigest()[:16]
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert to dictionary."""
|
|
data = asdict(self)
|
|
data['id'] = self.id
|
|
data['type'] = self.type.value
|
|
data['status'] = self.status.value
|
|
return data
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Dict[str, Any]) -> Capability:
|
|
"""Create from dictionary."""
|
|
data = data.copy()
|
|
data.pop('id', None)
|
|
return cls(**data)
|
|
|
|
|
|
@dataclass
|
|
class CapabilityQuery:
|
|
"""Query for searching capabilities."""
|
|
name_pattern: Optional[str] = None
|
|
types: Optional[List[CapabilityType]] = None
|
|
tags: Optional[List[str]] = None
|
|
status: Optional[CapabilityStatus] = None
|
|
min_version: Optional[str] = None
|
|
required_deps: Optional[List[str]] = None
|
|
|
|
|
|
class CapabilityRegistry:
|
|
"""Registry for managing agent capabilities."""
|
|
|
|
def __init__(self, agent_id: str):
|
|
self.agent_id = agent_id
|
|
self._capabilities: Dict[str, Capability] = {}
|
|
self._handlers: Dict[str, Callable] = {}
|
|
self._lock = threading.RLock()
|
|
self._version_counter = 0
|
|
|
|
def register(self, capability: Capability, handler: Optional[Callable] = None) -> str:
|
|
"""
|
|
Register a new capability.
|
|
|
|
Args:
|
|
capability: The capability to register
|
|
handler: Optional handler function for this capability
|
|
|
|
Returns:
|
|
Capability ID
|
|
"""
|
|
with self._lock:
|
|
cap_id = capability.id
|
|
|
|
# Update timestamp if already exists
|
|
if cap_id in self._capabilities:
|
|
capability.updated_at = datetime.utcnow().isoformat()
|
|
|
|
self._capabilities[cap_id] = capability
|
|
|
|
if handler:
|
|
self._handlers[cap_id] = handler
|
|
|
|
self._version_counter += 1
|
|
|
|
return cap_id
|
|
|
|
def unregister(self, cap_id: str) -> bool:
|
|
"""Unregister a capability."""
|
|
with self._lock:
|
|
if cap_id in self._capabilities:
|
|
del self._capabilities[cap_id]
|
|
self._handlers.pop(cap_id, None)
|
|
self._version_counter += 1
|
|
return True
|
|
return False
|
|
|
|
def get(self, cap_id: str) -> Optional[Capability]:
|
|
"""Get a capability by ID."""
|
|
with self._lock:
|
|
return self._capabilities.get(cap_id)
|
|
|
|
def get_by_name(self, name: str, version: Optional[str] = None) -> Optional[Capability]:
|
|
"""Get a capability by name and optional version."""
|
|
with self._lock:
|
|
for cap in self._capabilities.values():
|
|
if cap.name == name:
|
|
if version is None or cap.version == version:
|
|
return cap
|
|
return None
|
|
|
|
def list_all(self) -> List[Capability]:
|
|
"""List all registered capabilities."""
|
|
with self._lock:
|
|
return list(self._capabilities.values())
|
|
|
|
def query(self, query: CapabilityQuery) -> List[Capability]:
|
|
"""
|
|
Query capabilities with filters.
|
|
|
|
Args:
|
|
query: Query parameters
|
|
|
|
Returns:
|
|
List of matching capabilities
|
|
"""
|
|
with self._lock:
|
|
results = []
|
|
|
|
for cap in self._capabilities.values():
|
|
# Name pattern match
|
|
if query.name_pattern and query.name_pattern not in cap.name:
|
|
continue
|
|
|
|
# Type filter
|
|
if query.types and cap.type not in query.types:
|
|
continue
|
|
|
|
# Tags filter (match any)
|
|
if query.tags and not any(tag in cap.tags for tag in query.tags):
|
|
continue
|
|
|
|
# Status filter
|
|
if query.status and cap.status != query.status:
|
|
continue
|
|
|
|
# Dependencies check
|
|
if query.required_deps:
|
|
if not all(dep in cap.dependencies for dep in query.required_deps):
|
|
continue
|
|
|
|
results.append(cap)
|
|
|
|
return results
|
|
|
|
def get_handler(self, cap_id: str) -> Optional[Callable]:
|
|
"""Get the handler for a capability."""
|
|
with self._lock:
|
|
return self._handlers.get(cap_id)
|
|
|
|
def execute(self, cap_id: str, **kwargs) -> Any:
|
|
"""
|
|
Execute a capability handler.
|
|
|
|
Args:
|
|
cap_id: Capability ID
|
|
**kwargs: Arguments for the handler
|
|
|
|
Returns:
|
|
Handler result
|
|
"""
|
|
handler = self.get_handler(cap_id)
|
|
if handler is None:
|
|
raise ValueError(f"No handler registered for capability {cap_id}")
|
|
|
|
return handler(**kwargs)
|
|
|
|
def get_dependencies(self, cap_id: str) -> List[Capability]:
|
|
"""Get all dependencies for a capability."""
|
|
with self._lock:
|
|
cap = self._capabilities.get(cap_id)
|
|
if not cap:
|
|
return []
|
|
|
|
deps = []
|
|
for dep_name in cap.dependencies:
|
|
dep_cap = self.get_by_name(dep_name)
|
|
if dep_cap:
|
|
deps.append(dep_cap)
|
|
|
|
return deps
|
|
|
|
def check_dependencies(self, cap_id: str) -> Dict[str, bool]:
|
|
"""
|
|
Check if all dependencies are satisfied.
|
|
|
|
Returns:
|
|
Dict mapping dependency names to satisfaction status
|
|
"""
|
|
with self._lock:
|
|
cap = self._capabilities.get(cap_id)
|
|
if not cap:
|
|
return {}
|
|
|
|
return {
|
|
dep: self.get_by_name(dep) is not None
|
|
for dep in cap.dependencies
|
|
}
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
"""Get registry statistics."""
|
|
with self._lock:
|
|
type_counts = {}
|
|
status_counts = {}
|
|
|
|
for cap in self._capabilities.values():
|
|
type_counts[cap.type.value] = type_counts.get(cap.type.value, 0) + 1
|
|
status_counts[cap.status.value] = status_counts.get(cap.status.value, 0) + 1
|
|
|
|
return {
|
|
"total_capabilities": len(self._capabilities),
|
|
"version": self._version_counter,
|
|
"by_type": type_counts,
|
|
"by_status": status_counts,
|
|
"with_handlers": len(self._handlers)
|
|
}
|
|
|
|
def export(self) -> Dict[str, Any]:
|
|
"""Export all capabilities as dictionary."""
|
|
with self._lock:
|
|
return {
|
|
"agent_id": self.agent_id,
|
|
"version": self._version_counter,
|
|
"exported_at": datetime.utcnow().isoformat(),
|
|
"capabilities": [cap.to_dict() for cap in self._capabilities.values()]
|
|
}
|
|
|
|
def export_json(self, indent: int = 2) -> str:
|
|
"""Export all capabilities as JSON string."""
|
|
return json.dumps(self.export(), indent=indent)
|
|
|
|
def import_capabilities(self, data: Dict[str, Any]) -> int:
|
|
"""Import capabilities from dictionary."""
|
|
with self._lock:
|
|
count = 0
|
|
for cap_data in data.get("capabilities", []):
|
|
cap = Capability.from_dict(cap_data)
|
|
self.register(cap)
|
|
count += 1
|
|
return count
|
|
|
|
|
|
# Singleton registry instance
|
|
_registry_instance: Optional[CapabilityRegistry] = None
|
|
|
|
|
|
def get_registry(agent_id: Optional[str] = None) -> CapabilityRegistry:
|
|
"""Get or create the global registry instance."""
|
|
global _registry_instance
|
|
if _registry_instance is None:
|
|
if agent_id is None:
|
|
agent_id = "ap-agent"
|
|
_registry_instance = CapabilityRegistry(agent_id)
|
|
return _registry_instance
|
|
|
|
|
|
def reset_registry():
|
|
"""Reset the global registry (mainly for testing)."""
|
|
global _registry_instance
|
|
_registry_instance = None
|
|
|
|
|
|
# Convenience decorators
|
|
def capability(
|
|
name: str,
|
|
cap_type: CapabilityType = CapabilityType.SKILL,
|
|
version: str = "1.0.0",
|
|
description: str = "",
|
|
**kwargs
|
|
):
|
|
"""
|
|
Decorator to register a function as a capability.
|
|
|
|
Usage:
|
|
@capability("my_tool", CapabilityType.TOOL, "1.0.0", "Does something")
|
|
def my_tool_function(arg1: str) -> str:
|
|
return f"Result: {arg1}"
|
|
"""
|
|
def decorator(func: Callable) -> Callable:
|
|
cap = Capability(
|
|
name=name,
|
|
type=cap_type,
|
|
version=version,
|
|
description=description,
|
|
**kwargs
|
|
)
|
|
get_registry().register(cap, handler=func)
|
|
return func
|
|
return decorator
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Demo usage
|
|
registry = CapabilityRegistry("test-agent")
|
|
|
|
# Register some capabilities
|
|
cap1 = Capability(
|
|
name="file_reader",
|
|
type=CapabilityType.TOOL,
|
|
version="1.0.0",
|
|
description="Read file contents",
|
|
parameters={"path": {"type": "string", "required": True}},
|
|
tags=["filesystem", "io"]
|
|
)
|
|
|
|
cap2 = Capability(
|
|
name="code_analyzer",
|
|
type=CapabilityType.SKILL,
|
|
version="2.1.0",
|
|
description="Analyze code for issues",
|
|
dependencies=["file_reader"],
|
|
tags=["code", "analysis"]
|
|
)
|
|
|
|
registry.register(cap1)
|
|
registry.register(cap2)
|
|
|
|
print("Registry stats:", registry.get_stats())
|
|
print("\nAll capabilities:")
|
|
for cap in registry.list_all():
|
|
print(f" - {cap.name} ({cap.version}): {cap.description}")
|
|
|
|
# Query example
|
|
query = CapabilityQuery(tags=["code"])
|
|
results = registry.query(query)
|
|
print(f"\nQuery results for 'code' tag: {[r.name for r in results]}")
|
|
|
|
# Export
|
|
print("\nExport:")
|
|
print(registry.export_json())
|