evotraders/backend/agents/team/registry.py

# -*- coding: utf-8 -*-
"""AgentRegistry - Agent registration and lookup by role.

Provides register(), unregister(), and get_by_role() for agent
discovery and management.
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

from agentscope.message import Msg

logger = logging.getLogger(__name__)


class AgentRegistry:
    """Registry for agent instances with role-based lookup.

    Supports:
    - register(): Add agent with roles
    - unregister(): Remove agent
    - get_by_role(): Find agents by role
    - get_by_id(): Get specific agent

    Each agent can have multiple roles for flexible dispatch.
    """

    def __init__(self):
        self._agents: Dict[str, Any] = {}
        self._roles: Dict[str, List[str]] = {}
        self._agent_roles: Dict[str, List[str]] = {}

    def register(
        self,
        agent_id: str,
        agent: Any,
        roles: Optional[List[str]] = None,
    ) -> None:
        """Register an agent with optional roles.

        Args:
            agent_id: Unique agent identifier
            agent: Agent instance
            roles: Optional list of role strings
        """
        self._agents[agent_id] = agent
        self._agent_roles[agent_id] = roles or []

        for role in self._agent_roles[agent_id]:
            if role not in self._roles:
                self._roles[role] = []
            if agent_id not in self._roles[role]:
                self._roles[role].append(agent_id)

        logger.info(
            "Registered agent %s with roles %s",
            agent_id,
            self._agent_roles[agent_id],
        )

    def unregister(self, agent_id: str) -> bool:
        """Unregister an agent.

        Args:
            agent_id: Agent identifier to remove

        Returns:
            True if agent was removed
        """
        if agent_id not in self._agents:
            return False

        roles = self._agent_roles.pop(agent_id, [])
        for role in roles:
            if role in self._roles:
                try:
                    self._roles[role].remove(agent_id)
                except ValueError:
                    pass

        del self._agents[agent_id]
        logger.info("Unregistered agent: %s", agent_id)
        return True

    def get_by_id(self, agent_id: str) -> Optional[Any]:
        """Get agent by ID.

        Args:
            agent_id: Agent identifier

        Returns:
            Agent instance or None
        """
        return self._agents.get(agent_id)

    def get_by_role(self, role: str) -> List[Any]:
        """Get all agents with a given role.

        Args:
            role: Role string to search for

        Returns:
            List of agent instances with the role
        """
        agent_ids = self._roles.get(role, [])
        return [self._agents[aid] for aid in agent_ids if aid in self._agents]

    def get_by_roles(self, roles: List[str]) -> List[Any]:
        """Get agents matching ANY of the given roles.

        Args:
            roles: List of role strings

        Returns:
            List of unique agent instances matching any role
        """
        seen = set()
        result = []
        for role in roles:
            for agent in self.get_by_role(role):
                if id(agent) not in seen:
                    seen.add(id(agent))
                    result.append(agent)
        return result

    def list_agents(self) -> List[str]:
        """List all registered agent IDs.

        Returns:
            List of agent identifiers
        """
        return list(self._agents.keys())

    def list_roles(self) -> List[str]:
        """List all registered roles.

        Returns:
            List of role strings
        """
        return list(self._roles.keys())

    def list_roles_for_agent(self, agent_id: str) -> List[str]:
        """List roles for specific agent.

        Args:
            agent_id: Agent identifier

        Returns:
            List of role strings
        """
        return list(self._agent_roles.get(agent_id, []))

    def update_roles(self, agent_id: str, roles: List[str]) -> None:
        """Update roles for an existing agent.

        Args:
            agent_id: Agent identifier
            roles: New list of roles
        """
        if agent_id not in self._agents:
            raise KeyError(f"Agent not registered: {agent_id}")

        old_roles = self._agent_roles.get(agent_id, [])
        for role in old_roles:
            if role in self._roles:
                try:
                    self._roles[role].remove(agent_id)
                except ValueError:
                    pass

        self._agent_roles[agent_id] = roles
        for role in roles:
            if role not in self._roles:
                self._roles[role] = []
            if agent_id not in self._roles[role]:
                self._roles[role].append(agent_id)

        logger.info("Updated roles for agent %s: %s", agent_id, roles)

    @property
    def agents(self) -> Dict[str, Any]:
        """Get copy of registered agents dict."""
        return dict(self._agents)


__all__ = ["AgentRegistry"]