catonline_ai/vw-agentic-rag/service/utils/error_handler.py

"""
DRY Error Handling and Logging Utilities
"""

import json
import logging
import traceback
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Dict, Optional, Callable
from functools import wraps

from ..sse import create_error_event, create_tool_error_event


class ErrorCode(Enum):
    """Error codes for different types of failures"""
    # Client errors (4xxx)
    INVALID_REQUEST = 4001
    MISSING_PARAMETERS = 4002
    INVALID_SESSION = 4003
    
    # Server errors (5xxx)
    LLM_ERROR = 5001
    TOOL_ERROR = 5002
    DATABASE_ERROR = 5003
    MEMORY_ERROR = 5004
    EXTERNAL_API_ERROR = 5005
    INTERNAL_ERROR = 5000


class ErrorCategory(Enum):
    """Error categories for better organization"""
    VALIDATION = "validation"
    LLM = "llm"
    TOOL = "tool"
    DATABASE = "database"
    MEMORY = "memory"
    EXTERNAL_API = "external_api"
    INTERNAL = "internal"


class StructuredLogger:
    """DRY structured logging with automatic error handling"""
    
    def __init__(self, name: str):
        self.logger = logging.getLogger(name)
    
    def error(self, msg: str, error: Optional[Exception] = None, category: ErrorCategory = ErrorCategory.INTERNAL, 
              error_code: ErrorCode = ErrorCode.INTERNAL_ERROR, extra: Optional[Dict[str, Any]] = None):
        """Log structured error with stack trace"""
        data: Dict[str, Any] = {
            "message": msg,
            "category": category.value,
            "error_code": error_code.value,
            "timestamp": datetime.now(timezone.utc).isoformat()
        }
        
        if error:
            data.update({
                "error_type": type(error).__name__,
                "error_message": str(error),
                "stack_trace": traceback.format_exc()
            })
        
        if extra:
            data["extra"] = extra
            
        self.logger.error(json.dumps(data))
    
    def info(self, msg: str, extra: Optional[Dict[str, Any]] = None):
        """Log structured info"""
        data: Dict[str, Any] = {"message": msg, "timestamp": datetime.now(timezone.utc).isoformat()}
        if extra:
            data["extra"] = extra
        self.logger.info(json.dumps(data))
    
    def warning(self, msg: str, extra: Optional[Dict[str, Any]] = None):
        """Log structured warning"""
        data: Dict[str, Any] = {"message": msg, "timestamp": datetime.now(timezone.utc).isoformat()}
        if extra:
            data["extra"] = extra
        self.logger.warning(json.dumps(data))


def get_user_message(category: ErrorCategory) -> str:
    """Get user-friendly error messages in English"""
    messages = {
        ErrorCategory.VALIDATION: "Invalid request parameters. Please check your input.",
        ErrorCategory.LLM: "AI service is temporarily unavailable. Please try again later.",
        ErrorCategory.TOOL: "Tool execution failed. Please retry your request.",
        ErrorCategory.DATABASE: "Database service is temporarily unavailable.",
        ErrorCategory.MEMORY: "Session storage issue occurred. Please refresh the page.",
        ErrorCategory.EXTERNAL_API: "External service connection failed.",
        ErrorCategory.INTERNAL: "Internal server error. We are working to resolve this."
    }
    return messages.get(category, "Unknown error occurred. Please contact technical support.")


def handle_async_errors(category: ErrorCategory, error_code: ErrorCode, 
                       stream_callback: Optional[Callable] = None, tool_id: Optional[str] = None):
    """DRY decorator for async error handling with streaming support"""
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            logger = StructuredLogger(func.__module__)
            
            try:
                return await func(*args, **kwargs)
            except Exception as e:
                user_msg = get_user_message(category)
                
                logger.error(
                    f"Error in {func.__name__}: {str(e)}",
                    error=e,
                    category=category,
                    error_code=error_code,
                    extra={"function": func.__name__, "args_count": len(args)}
                )
                
                # Send error event if streaming
                if stream_callback:
                    if tool_id:
                        await stream_callback(create_tool_error_event(tool_id, func.__name__, user_msg))
                    else:
                        await stream_callback(create_error_event(user_msg))
                
                # Re-raise with user-friendly message for API responses
                raise Exception(user_msg) from e
        return wrapper
    return decorator


def handle_sync_errors(category: ErrorCategory, error_code: ErrorCode):
    """DRY decorator for sync error handling"""
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            logger = StructuredLogger(func.__module__)
            
            try:
                return func(*args, **kwargs)
            except Exception as e:
                logger.error(
                    f"Error in {func.__name__}: {str(e)}",
                    error=e,
                    category=category,
                    error_code=error_code,
                    extra={"function": func.__name__}
                )
                raise Exception(get_user_message(category)) from e
        return wrapper
    return decorator


def create_error_response(category: ErrorCategory, error_code: ErrorCode, 
                         technical_msg: Optional[str] = None) -> Dict[str, Any]:
    """Create consistent error response format"""
    return {
        "user_message": get_user_message(category),
        "error_code": error_code.value,
        "category": category.value,
        "technical_message": technical_msg,
        "timestamp": datetime.now(timezone.utc).isoformat()
    }