Update utils/error_handling.py

utils/error_handling.py

@@ -1,147 +1,256 @@
-import traceback
+"""
+Error handling utilities for the MONA application
+Fixed version with proper exception handling and logging
+"""
+
 import functools
-from typing import Any, Callable
+import traceback
+from typing import Any, Callable, Dict, Optional, Union
+
+from utils.logging import get_logger, log_error, log_warning
 
-# Import your logging utilities
-from utils.logging import get_logger, log_error
 
-# Custom exception classes
 class DataError(Exception):
     """Custom exception for data-related errors"""
-    pass
+    def __init__(self, message: str, error_code: str = None, details: Dict = None):
+        self.message = message
+        self.error_code = error_code or "DATA_ERROR"
+        self.details = details or {}
+        super().__init__(self.message)
+
 
 class ValidationError(Exception):
     """Custom exception for validation errors"""
-    pass
+    def __init__(self, message: str, field: str = None, value: Any = None):
+        self.message = message
+        self.field = field
+        self.value = value
+        super().__init__(self.message)
 
-class AIModelError(Exception):
-    """Custom exception for AI model errors"""
-    def __init__(self, message, metadata=None):
-        super().__init__(message)
-        self.metadata = metadata or {}
 
-# Get logger
-logger = get_logger(__name__)
+class ProcessingError(Exception):
+    """Custom exception for processing errors"""
+    def __init__(self, message: str, step: str = None, context: Dict = None):
+        self.message = message
+        self.step = step
+        self.context = context or {}
+        super().__init__(self.message)
+
 
 def handle_data_exceptions(func: Callable) -> Callable:
     """
-    Decorator to handle data-related exceptions gracefully
+    Decorator to handle data-related exceptions
+    
+    Args:
+        func: Function to wrap with exception handling
+        
+    Returns:
+        Wrapped function with exception handling
     """
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
         try:
             return func(*args, **kwargs)
-        except (DataError, ValidationError) as e:
-            log_error(logger, f"Data handling error in {func.__name__}: {str(e)}")
-            print(f"Data Error: {str(e)}")
-            return None
+        except DataError as e:
+            log_error(f"Data error in {func.__name__}: {e.message}", 
+                     error=e, 
+                     extra_data={"error_code": e.error_code, "details": e.details})
+            raise
+        except ValidationError as e:
+            log_error(f"Validation error in {func.__name__}: {e.message}",
+                     error=e,
+                     extra_data={"field": e.field, "value": e.value})
+            raise
+        except ProcessingError as e:
+            log_error(f"Processing error in {func.__name__}: {e.message}",
+                     error=e,
+                     extra_data={"step": e.step, "context": e.context})
+            raise
         except Exception as e:
-            log_error(logger, f"Unexpected error in {func.__name__}: {str(e)}")
-            log_error(logger, f"Traceback: {traceback.format_exc()}")
-            print(f"An unexpected error occurred: {str(e)}")
-            return None
+            log_error(f"Unexpected error in {func.__name__}: {str(e)}", 
+                     error=e)
+            raise ProcessingError(f"Unexpected error in {func.__name__}", 
+                                context={"original_error": str(e)})
+    
     return wrapper
 
-def handle_ai_model_exceptions(func: Callable) -> Callable:
-    """
-    Decorator to handle AI model-related exceptions gracefully
-    """
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        try:
-            return func(*args, **kwargs)
-        except (AIModelError, ValidationError) as e:
-            log_error(logger, f"AI model error in {func.__name__}: {str(e)}")
-            print(f"AI Model Error: {str(e)}")
-            return None
-        except Exception as e:
-            log_error(logger, f"Unexpected AI model error in {func.__name__}: {str(e)}")
-            log_error(logger, f"Traceback: {traceback.format_exc()}")
-            print(f"An unexpected AI model error occurred: {str(e)}")
-            return None
-    return wrapper
 
-def handle_exceptions(func: Callable) -> Callable:
+def safe_execute(func: Callable, *args, default_return=None, **kwargs) -> tuple:
     """
-    General exception handler decorator
+    Safely execute a function and return success status with result
+    
+    Args:
+        func: Function to execute
+        *args: Positional arguments for the function
+        default_return: Default value to return on error
+        **kwargs: Keyword arguments for the function
+        
+    Returns:
+        tuple: (success: bool, result: Any, error: Optional[Exception])
     """
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        try:
-            return func(*args, **kwargs)
-        except Exception as e:
-            log_error(logger, f"Exception in {func.__name__}: {str(e)}")
-            log_error(logger, f"Traceback: {traceback.format_exc()}")
-            return None
-    return wrapper
+    try:
+        result = func(*args, **kwargs)
+        return True, result, None
+    except Exception as e:
+        log_error(f"Error executing {func.__name__}: {str(e)}", error=e)
+        return False, default_return, e
 
-def validate_data(data: Any, data_type: str) -> bool:
+
+def validate_data(data: Any, validation_rules: Dict) -> bool:
     """
-    Validate data based on expected type
+    Validate data against a set of rules
+    
+    Args:
+        data: Data to validate
+        validation_rules: Dictionary of validation rules
+        
+    Returns:
+        bool: True if validation passes
+        
+    Raises:
+        ValidationError: If validation fails
     """
-    if data is None:
-        raise ValidationError(f"Data cannot be None for type: {data_type}")
-
-    if data_type == "dataframe":
-        try:
-            import pandas as pd
-            if not isinstance(data, pd.DataFrame):
-                raise ValidationError("Expected pandas DataFrame")
-            if data.empty:
-                raise ValidationError("DataFrame cannot be empty")
-        except ImportError:
-            raise ValidationError("pandas not available for DataFrame validation")
-
-    elif data_type == "list":
-        if not isinstance(data, list):
-            raise ValidationError("Expected list")
-        if len(data) == 0:
-            raise ValidationError("List cannot be empty")
-
-    elif data_type == "dict":
-        if not isinstance(data, dict):
-            raise ValidationError("Expected dictionary")
-        if len(data) == 0:
-            raise ValidationError("Dictionary cannot be empty")
-
+    if not isinstance(validation_rules, dict):
+        raise ValidationError("Validation rules must be a dictionary")
+    
+    for field, rules in validation_rules.items():
+        if isinstance(data, dict):
+            value = data.get(field)
+        else:
+            value = getattr(data, field, None)
+        
+        # Check required fields
+        if rules.get('required', False) and value is None:
+            raise ValidationError(f"Field '{field}' is required", field=field, value=value)
+        
+        # Check data type
+        if 'type' in rules and value is not None:
+            expected_type = rules['type']
+            if not isinstance(value, expected_type):
+                raise ValidationError(
+                    f"Field '{field}' must be of type {expected_type.__name__}, got {type(value).__name__}",
+                    field=field, 
+                    value=value
+                )
+        
+        # Check value range for numbers
+        if 'min_value' in rules and isinstance(value, (int, float)):
+            if value < rules['min_value']:
+                raise ValidationError(
+                    f"Field '{field}' must be >= {rules['min_value']}, got {value}",
+                    field=field,
+                    value=value
+                )
+        
+        if 'max_value' in rules and isinstance(value, (int, float)):
+            if value > rules['max_value']:
+                raise ValidationError(
+                    f"Field '{field}' must be <= {rules['max_value']}, got {value}",
+                    field=field,
+                    value=value
+                )
+        
+        # Check string length
+        if 'min_length' in rules and isinstance(value, str):
+            if len(value) < rules['min_length']:
+                raise ValidationError(
+                    f"Field '{field}' must be at least {rules['min_length']} characters",
+                    field=field,
+                    value=value
+                )
+        
+        if 'max_length' in rules and isinstance(value, str):
+            if len(value) > rules['max_length']:
+                raise ValidationError(
+                    f"Field '{field}' must be at most {rules['max_length']} characters",
+                    field=field,
+                    value=value
+                )
+    
     return True
 
-def show_error(message: str):
-    """
-    Display error message (streamlit-safe)
-    """
-    try:
-        import streamlit as st
-        st.error(message)
-    except ImportError:
-        print(f"ERROR: {message}")
 
-def show_warning(message: str):
+def format_error_response(error: Exception) -> Dict:
     """
-    Display warning message (streamlit-safe)
+    Format an error into a standardized response
+    
+    Args:
+        error: Exception to format
+        
+    Returns:
+        Dict: Formatted error response
     """
-    try:
-        import streamlit as st
-        st.warning(message)
-    except ImportError:
-        print(f"WARNING: {message}")
+    if isinstance(error, (DataError, ValidationError, ProcessingError)):
+        return {
+            "success": False,
+            "error_type": error.__class__.__name__,
+            "message": error.message,
+            "details": getattr(error, 'details', {}),
+            "field": getattr(error, 'field', None),
+            "error_code": getattr(error, 'error_code', None)
+        }
+    else:
+        return {
+            "success": False,
+            "error_type": "UnexpectedError",
+            "message": str(error),
+            "details": {}
+        }
 
-def show_success(message: str):
-    """
-    Display success message (streamlit-safe)
-    """
-    try:
-        import streamlit as st
-        st.success(message)
-    except ImportError:
-        print(f"SUCCESS: {message}")
 
-def safe_get(dictionary: dict, key: str, default: Any = None) -> Any:
+def retry_on_failure(max_retries: int = 3, delay: float = 1.0, exceptions: tuple = (Exception,)):
     """
-    Safely get value from dictionary
+    Decorator to retry function execution on failure
+    
+    Args:
+        max_retries: Maximum number of retry attempts
+        delay: Delay between retries in seconds
+        exceptions: Tuple of exceptions to catch and retry on
+        
+    Returns:
+        Decorator function
     """
-    try:
-        return dictionary.get(key, default)
-    except (AttributeError, TypeError):
-        logger.warning(f"safe_get failed for key '{key}' - dictionary is not dict-like")
-        return default
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            import time
+            
+            last_exception = None
+            for attempt in range(max_retries + 1):
+                try:
+                    return func(*args, **kwargs)
+                except exceptions as e:
+                    last_exception = e
+                    if attempt < max_retries:
+                        log_warning(f"Attempt {attempt + 1} failed for {func.__name__}: {str(e)}. Retrying in {delay}s...")
+                        time.sleep(delay)
+                    else:
+                        log_error(f"All {max_retries + 1} attempts failed for {func.__name__}", error=e)
+            
+            raise last_exception
+        
+        return wrapper
+    return decorator
+
+
+# Context manager for error handling
+class ErrorHandler:
+    """Context manager for comprehensive error handling"""
+    
+    def __init__(self, operation_name: str, raise_on_error: bool = True):
+        self.operation_name = operation_name
+        self.raise_on_error = raise_on_error
+        self.error = None
+        
+    def __enter__(self):
+        return self
+        
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is not None:
+            self.error = exc_val
+            log_error(f"Error in {self.operation_name}: {str(exc_val)}", error=exc_val)
+            
+            if not self.raise_on_error:
+                return True  # Suppress the exception
+        
+        return False  # Let the exception propagate