visualization/agents/base_agent.py

"""
Base Agent class for visualization agents
Provides common functionality and interface for all agents
"""
from abc import ABC, abstractmethod
from typing import Dict, Any
import logging


class BaseVisualizationAgent(ABC):
    """
    Abstract base class for all visualization agents
    """

    def __init__(self, name: str, model: str = "gpt-5-nano", temperature: float = 0.7):
        """
        Initialize base agent

        Args:
            name: Agent name
            model: LLM model to use
            temperature: LLM temperature
        """
        self.name = name
        self.model = model
        self.temperature = temperature
        self.logger = logging.getLogger(f"visualization.agents.{name}")

    @abstractmethod
    def process(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process input data and return results

        Args:
            input_data: Input data dictionary

        Returns:
            Results dictionary
        """
        pass

    @abstractmethod
    def validate_input(self, input_data: Dict[str, Any]) -> bool:
        """
        Validate input data

        Args:
            input_data: Input data dictionary

        Returns:
            True if valid, False otherwise
        """
        pass

    def log_processing(self, message: str, level: str = "info"):
        """
        Log processing information

        Args:
            message: Log message
            level: Log level (info, warning, error)
        """
        log_func = getattr(self.logger, level.lower(), self.logger.info)
        log_func(f"[{self.name}] {message}")

    def handle_error(self, error: Exception, context: str = "") -> Dict[str, Any]:
        """
        Handle errors consistently

        Args:
            error: Exception that occurred
            context: Additional context information

        Returns:
            Error response dictionary
        """
        error_msg = f"Error in {self.name}: {str(error)}"
        if context:
            error_msg += f" | Context: {context}"

        self.log_processing(error_msg, level="error")

        return {
            'success': False,
            'error': str(error),
            'error_type': type(error).__name__,
            'context': context
        }