feat(agent-os): add agent registry runtime implementation

This commit is contained in:
2026-07-07 10:15:05 +08:00
parent 8a5ce63069
commit a5fdfbf057
+20 -100
View File
@@ -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