services/agent_registry.py

"""Registry for managing multiple LangGraph agents."""
from typing import Dict, Callable, Any
from langchain_core.language_models.chat_models import BaseChatModel
from domain.enums import AgentType
from graphs.base_graph import create_simple_graph, create_simple_graph_with_history


class AgentRegistry:
    """
    Registry for managing multiple agent graph builders.
    
    This allows for easy addition of new agent types without modifying
    the API layer. Each agent type maps to a graph builder function.
    """
    
    def __init__(self):
        """Initialize the agent registry with default agents."""
        self._builders: Dict[AgentType, Callable[[BaseChatModel], Any]] = {
            AgentType.SIMPLE: create_simple_graph,
            # AgentType.RAG: create_rag_graph,  # À implémenter plus tard
            # AgentType.TOOLS: create_tools_graph,  # À implémenter plus tard
        }
        
        self._descriptions = {
            AgentType.SIMPLE: "Simple conversational agent without tools or memory",
            AgentType.RAG: "Agent with Retrieval Augmented Generation (not yet implemented)",
            AgentType.TOOLS: "Agent with tools like web search, calculator (not yet implemented)",
            AgentType.CUSTOM: "Custom agent graph (not yet implemented)"
        }
    
    def register_agent(
        self,
        agent_type: AgentType,
        builder: Callable[[BaseChatModel], Any],
        description: str = ""
    ) -> None:
        """
        Register a new agent builder.
        
        Args:
            agent_type: Type of agent
            builder: Function that takes an LLM and returns a compiled graph
            description: Description of the agent
        """
        self._builders[agent_type] = builder
        if description:
            self._descriptions[agent_type] = description
    
    def get_builder(self, agent_type: AgentType) -> Callable[[BaseChatModel], Any]:
        """
        Get the builder function for an agent type.
        
        Args:
            agent_type: Type of agent
            
        Returns:
            Builder function
            
        Raises:
            ValueError: If agent type is not registered
        """
        if agent_type not in self._builders:
            raise ValueError(
                f"Agent type '{agent_type}' not implemented. "
                f"Available types: {list(self._builders.keys())}"
            )
        return self._builders[agent_type]
    
    def is_available(self, agent_type: AgentType) -> bool:
        """
        Check if an agent type is available.
        
        Args:
            agent_type: Type of agent
            
        Returns:
            True if agent is available, False otherwise
        """
        return agent_type in self._builders
    
    def list_agents(self) -> list[dict]:
        """
        List all registered agents with their information.
        
        Returns:
            List of agent information dictionaries
        """
        agents = []
        for agent_type in AgentType:
            agents.append({
                "type": agent_type.value,
                "name": agent_type.value.capitalize(),
                "description": self._descriptions.get(
                    agent_type,
                    "No description available"
                ),
                "available": self.is_available(agent_type)
            })
        return agents


# Singleton instance
agent_registry = AgentRegistry()