Update utils/error_handling.py

utils/error_handling.py

@@ -1,6 +1,6 @@
 import traceback
-from typing import Any, Callable
 import functools
+from typing import Any, Callable
 
 # Import your logging utilities
 from utils.logging import get_logger, log_error
@@ -14,18 +14,18 @@ class ValidationError(Exception):
     """Custom exception for validation errors"""
     pass
 
+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__)
 
 def handle_data_exceptions(func: Callable) -> Callable:
     """
     Decorator to handle data-related exceptions gracefully
-    
-    Args:
-        func: Function to wrap with exception handling
-    
-    Returns:
-        Wrapped function with exception handling
     """
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
@@ -33,7 +33,6 @@ def handle_data_exceptions(func: Callable) -> Callable:
             return func(*args, **kwargs)
         except (DataError, ValidationError) as e:
             log_error(logger, f"Data handling error in {func.__name__}: {str(e)}")
-            # Instead of streamlit error, we'll raise the exception or return None
             print(f"Data Error: {str(e)}")
             return None
         except Exception as e:
@@ -41,26 +40,48 @@ def handle_data_exceptions(func: Callable) -> Callable:
             log_error(logger, f"Traceback: {traceback.format_exc()}")
             print(f"An unexpected error occurred: {str(e)}")
             return None
-    
+    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:
+    """
+    General exception handler decorator
+    """
+    @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
 
 def validate_data(data: Any, data_type: str) -> bool:
     """
     Validate data based on expected type
-    
-    Args:
-        data: Data to validate
-        data_type: Expected data type
-    
-    Returns:
-        bool: True if valid, False otherwise
-    
-    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
@@ -70,19 +91,19 @@ def validate_data(data: Any, data_type: str) -> bool:
                 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")
-    
+
     return True
 
 def show_error(message: str):
@@ -115,41 +136,12 @@ def show_success(message: str):
     except ImportError:
         print(f"SUCCESS: {message}")
 
-def handle_exceptions(func: Callable) -> Callable:
-    """
-    General exception handler decorator
-    
-    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 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
-
 def safe_get(dictionary: dict, key: str, default: Any = None) -> Any:
     """
     Safely get value from dictionary
-    
-    Args:
-        dictionary: Dictionary to get value from
-        key: Key to look up
-        default: Default value if key not found
-    
-    Returns:
-        Value from dictionary or default
     """
     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
+        return default