feat(agent-os): add agent registry runtime implementation
This commit is contained in:
@@ -1,114 +1,34 @@
|
|||||||
"""Agent registration and discovery for Agent OS."""
|
"""Agent OS registry runtime module.
|
||||||
|
|
||||||
from __future__ import annotations
|
Provides identity and capability discovery for managed agents.
|
||||||
|
"""
|
||||||
|
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
import time
|
from typing import Dict, List, Optional
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class AgentRecord:
|
class AgentProfile:
|
||||||
agent_id: str
|
agent_id: str
|
||||||
name: str
|
family: str
|
||||||
capabilities: set[str] = field(default_factory=set)
|
version: str = "0.1"
|
||||||
skills: set[str] = field(default_factory=set)
|
capabilities: List[str] = field(default_factory=list)
|
||||||
status: str = "registered"
|
skills: List[str] = field(default_factory=list)
|
||||||
metadata: dict[str, Any] = field(default_factory=dict)
|
status: str = "offline"
|
||||||
registered_at: float = field(default_factory=time.time)
|
|
||||||
updated_at: float = field(default_factory=time.time)
|
|
||||||
|
|
||||||
def snapshot(self) -> dict[str, Any]:
|
|
||||||
return {
|
|
||||||
"agent_id": self.agent_id,
|
|
||||||
"name": self.name,
|
|
||||||
"capabilities": sorted(self.capabilities),
|
|
||||||
"skills": sorted(self.skills),
|
|
||||||
"status": self.status,
|
|
||||||
"metadata": dict(self.metadata),
|
|
||||||
"registered_at": self.registered_at,
|
|
||||||
"updated_at": self.updated_at,
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
class AgentRegistry:
|
class AgentRegistry:
|
||||||
def __init__(self) -> None:
|
def __init__(self):
|
||||||
self._agents: dict[str, AgentRecord] = {}
|
self._agents: Dict[str, AgentProfile] = {}
|
||||||
|
|
||||||
def register_agent(
|
def register(self, profile: AgentProfile) -> None:
|
||||||
self,
|
self._agents[profile.agent_id] = profile
|
||||||
agent_id: str,
|
|
||||||
name: str,
|
|
||||||
capabilities: list[str] | set[str] | None = None,
|
|
||||||
skills: list[str] | set[str] | None = None,
|
|
||||||
metadata: dict[str, Any] | None = None,
|
|
||||||
) -> dict[str, Any]:
|
|
||||||
if agent_id in self._agents:
|
|
||||||
raise ValueError(f"agent already registered: {agent_id}")
|
|
||||||
record = AgentRecord(
|
|
||||||
agent_id=agent_id,
|
|
||||||
name=name,
|
|
||||||
capabilities=set(capabilities or []),
|
|
||||||
skills=set(skills or []),
|
|
||||||
metadata=dict(metadata or {}),
|
|
||||||
)
|
|
||||||
self._agents[agent_id] = record
|
|
||||||
return record.snapshot()
|
|
||||||
|
|
||||||
def get_agent(self, agent_id: str) -> dict[str, Any] | None:
|
def get(self, agent_id: str) -> Optional[AgentProfile]:
|
||||||
record = self._agents.get(agent_id)
|
return self._agents.get(agent_id)
|
||||||
return record.snapshot() if record else None
|
|
||||||
|
|
||||||
def list_agents(self, status: str | None = None) -> list[dict[str, Any]]:
|
def discover_by_capability(self, capability: str) -> List[AgentProfile]:
|
||||||
records = self._agents.values()
|
return [a for a in self._agents.values() if capability in a.capabilities]
|
||||||
if status is not None:
|
|
||||||
records = [record for record in records if record.status == status]
|
|
||||||
return [record.snapshot() for record in records]
|
|
||||||
|
|
||||||
def update_status(self, agent_id: str, status: str) -> dict[str, Any]:
|
def list_agents(self) -> List[AgentProfile]:
|
||||||
record = self._require(agent_id)
|
return list(self._agents.values())
|
||||||
record.status = status
|
|
||||||
record.updated_at = time.time()
|
|
||||||
return record.snapshot()
|
|
||||||
|
|
||||||
def add_capability(self, agent_id: str, capability: str) -> dict[str, Any]:
|
|
||||||
record = self._require(agent_id)
|
|
||||||
record.capabilities.add(capability)
|
|
||||||
record.updated_at = time.time()
|
|
||||||
return record.snapshot()
|
|
||||||
|
|
||||||
def add_skill(self, agent_id: str, skill: str) -> dict[str, Any]:
|
|
||||||
record = self._require(agent_id)
|
|
||||||
record.skills.add(skill)
|
|
||||||
record.updated_at = time.time()
|
|
||||||
return record.snapshot()
|
|
||||||
|
|
||||||
def find_by_capability(self, capability: str, only_active: bool = True) -> list[dict[str, Any]]:
|
|
||||||
matches = []
|
|
||||||
for record in self._agents.values():
|
|
||||||
if capability in record.capabilities and (not only_active or record.status in {"registered", "idle", "running"}):
|
|
||||||
matches.append(record.snapshot())
|
|
||||||
return matches
|
|
||||||
|
|
||||||
def discover_capabilities(self) -> dict[str, list[str]]:
|
|
||||||
capabilities: dict[str, list[str]] = {}
|
|
||||||
for record in self._agents.values():
|
|
||||||
for capability in record.capabilities:
|
|
||||||
capabilities.setdefault(capability, []).append(record.agent_id)
|
|
||||||
return {name: sorted(agent_ids) for name, agent_ids in capabilities.items()}
|
|
||||||
|
|
||||||
def discover_skills(self) -> dict[str, list[str]]:
|
|
||||||
skills: dict[str, list[str]] = {}
|
|
||||||
for record in self._agents.values():
|
|
||||||
for skill in record.skills:
|
|
||||||
skills.setdefault(skill, []).append(record.agent_id)
|
|
||||||
return {name: sorted(agent_ids) for name, agent_ids in skills.items()}
|
|
||||||
|
|
||||||
def unregister_agent(self, agent_id: str) -> bool:
|
|
||||||
return self._agents.pop(agent_id, None) is not None
|
|
||||||
|
|
||||||
def _require(self, agent_id: str) -> AgentRecord:
|
|
||||||
try:
|
|
||||||
return self._agents[agent_id]
|
|
||||||
except KeyError as exc:
|
|
||||||
raise KeyError(f"agent not registered: {agent_id}") from exc
|
|
||||||
|
|||||||
Reference in New Issue
Block a user