from yera.agents.decorator import AgentFunctionWrapper


class AgentService:
    """Service for managing agents."""

    def __init__(self, agents: list[AgentFunctionWrapper]):
        """Initialise the agent service.

        Args:
            agents: The list of agents to manage.
        """
        self.agents: dict[str, AgentFunctionWrapper] = {
            agent.metadata.identifier: agent for agent in agents
        }

    def get_agent(self, agent_id: str) -> AgentFunctionWrapper:
        """Get an agent by ID.

        Args:
            agent_id: The ID of the agent to get.

        Returns:
            The agent function wrapper.

        Raises:
            ValueError: If the agent is not found.
        """
        try:
            return self.agents[agent_id]
        except KeyError:
            raise ValueError(f"Agent with ID '{agent_id}' not found") from None

    def list_agents(self) -> list[AgentFunctionWrapper]:
        """List all available agents.

        Returns:
            List of available agents.
        """
        return list(self.agents.values())
