gaia/utils/exceptions.py

#!/usr/bin/env python3
"""
Custom exception classes for the GAIA system.
"""

from typing import Optional, Any, Dict


class GAIAError(Exception):
    """Base exception for all GAIA-related errors."""
    
    def __init__(self, message: str, details: Optional[Dict[str, Any]] = None):
        super().__init__(message)
        self.message = message
        self.details = details or {}
    
    def __str__(self) -> str:
        if self.details:
            return f"{self.message} - Details: {self.details}"
        return self.message


class ModelError(GAIAError):
    """Exception raised for model-related errors."""
    
    def __init__(self, message: str, model_name: Optional[str] = None, 
                 provider: Optional[str] = None, **kwargs):
        super().__init__(message, kwargs)
        self.model_name = model_name
        self.provider = provider


class ModelNotAvailableError(ModelError):
    """Exception raised when requested model is not available."""
    pass


class ModelAuthenticationError(ModelError):
    """Exception raised for model authentication failures."""
    pass


class ModelOverloadedError(ModelError):
    """Exception raised when model is overloaded."""
    pass


class ToolError(GAIAError):
    """Exception raised for tool execution errors."""
    
    def __init__(self, message: str, tool_name: Optional[str] = None, 
                 input_data: Optional[Dict[str, Any]] = None, **kwargs):
        super().__init__(message, kwargs)
        self.tool_name = tool_name
        self.input_data = input_data


class ToolNotFoundError(ToolError):
    """Exception raised when requested tool is not found."""
    pass


class ToolValidationError(ToolError):
    """Exception raised for tool input validation errors."""
    pass


class ToolExecutionError(ToolError):
    """Exception raised during tool execution."""
    pass


class ToolTimeoutError(ToolError):
    """Exception raised when tool execution times out."""
    pass


class ClassificationError(GAIAError):
    """Exception raised for question classification errors."""
    
    def __init__(self, message: str, question: Optional[str] = None, **kwargs):
        super().__init__(message, kwargs)
        self.question = question


class FileProcessingError(GAIAError):
    """Exception raised for file processing errors."""
    
    def __init__(self, message: str, file_path: Optional[str] = None, 
                 file_type: Optional[str] = None, **kwargs):
        super().__init__(message, kwargs)
        self.file_path = file_path
        self.file_type = file_type


class APIError(GAIAError):
    """Exception raised for external API errors."""
    
    def __init__(self, message: str, api_name: Optional[str] = None,
                 status_code: Optional[int] = None, **kwargs):
        super().__init__(message, kwargs)
        self.api_name = api_name
        self.status_code = status_code


class ConfigurationError(GAIAError):
    """Exception raised for configuration errors."""
    pass


class ValidationError(GAIAError):
    """Exception raised for data validation errors."""
    
    def __init__(self, message: str, field: Optional[str] = None, 
                 value: Optional[Any] = None, **kwargs):
        super().__init__(message, kwargs)
        self.field = field
        self.value = value


# Error code mapping for consistent error handling
ERROR_CODES = {
    "MODEL_NOT_AVAILABLE": ModelNotAvailableError,
    "MODEL_AUTH_FAILED": ModelAuthenticationError,
    "MODEL_OVERLOADED": ModelOverloadedError,
    "TOOL_NOT_FOUND": ToolNotFoundError,
    "TOOL_VALIDATION_FAILED": ToolValidationError,
    "TOOL_EXECUTION_FAILED": ToolExecutionError,
    "TOOL_TIMEOUT": ToolTimeoutError,
    "CLASSIFICATION_FAILED": ClassificationError,
    "FILE_PROCESSING_FAILED": FileProcessingError,
    "API_ERROR": APIError,
    "CONFIG_ERROR": ConfigurationError,
    "VALIDATION_ERROR": ValidationError
}


def create_error(error_code: str, message: str, **kwargs) -> GAIAError:
    """Create error instance based on error code."""
    error_class = ERROR_CODES.get(error_code, GAIAError)
    return error_class(message, **kwargs)