services/agents/agent_registry.py

# services/agents/agent_registry.py
"""
Central registry for all utility agents.
Provides lazy initialization and retrieval.
"""
import os
from typing import Dict, Optional
from services.agents.base_agent import BaseUtilityAgent


class AgentRegistry:
    """
    Singleton registry for all utility agents.
    
    Lazy initialization ensures agents are only created when needed,
    avoiding startup overhead.
    """
    
    _instance: Optional['AgentRegistry'] = None
    _agents: Dict[str, BaseUtilityAgent] = {}
    _initialized: bool = False
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    def _initialize_agents(self):
        """Initialize all agents lazily."""
        if self._initialized:
            return
        
        # Import agents here to avoid circular imports
        # and to delay initialization until first use
        from services.agents.extract_text_agent import ExtractTextAgent
        from services.agents.extract_tables_agent import ExtractTablesAgent
        from services.agents.describe_images_agent import DescribeImagesAgent
        from services.agents.summarizer_agent import SummarizerAgent
        from services.agents.classifier_agent import ClassifierAgent
        from services.agents.ner_agent import NERAgent
        from services.agents.translator_agent import TranslatorAgent
        from services.agents.signature_verification_agent import SignatureVerificationAgent
        from services.agents.stamp_detection_agent import StampDetectionAgent
        
        # Register all agents
        self._agents = {
            "extract_text": ExtractTextAgent(),
            "extract_tables": ExtractTablesAgent(),
            "describe_images": DescribeImagesAgent(),
            "summarize": SummarizerAgent(),
            "classify": ClassifierAgent(),
            "ner": NERAgent(),
            "translate": TranslatorAgent(),
            "signature_verification": SignatureVerificationAgent(),
            "stamp_detection": StampDetectionAgent(),
        }
        
        self._initialized = True
    
    def get_agent(self, name: str) -> Optional[BaseUtilityAgent]:
        """
        Get agent by name.
        
        Args:
            name: Agent identifier (e.g., "extract_text")
            
        Returns:
            Agent instance or None if not found
        """
        if not self._initialized:
            self._initialize_agents()
        
        return self._agents.get(name)
    
    def list_agents(self) -> list:
        """Get list of all registered agent names."""
        if not self._initialized:
            self._initialize_agents()
        
        return list(self._agents.keys())


# Singleton instance
_registry = AgentRegistry()


def get_agent(name: str) -> Optional[BaseUtilityAgent]:
    """
    Convenience function to get agent from global registry.
    
    Args:
        name: Agent identifier
        
    Returns:
        Agent instance or None
    """
    return _registry.get_agent(name)


def use_agents_enabled() -> bool:
    """
    Check if agent mode is enabled via feature flag.
    
    Returns:
        True if USE_AGENTS=true in environment
    """
    return os.getenv("USE_AGENTS", "false").lower() == "true"