Features: - 🤖 AI image generation with Gemini models - 🎬 Video planning and storyboard creation - ✨ Advanced prompt optimization and analysis - 🔧 Smart API key configuration and management - 📊 Multiple model support (Gemini 2.5 Flash/Pro) - 🛡️ Robust error handling and rate limiting - 📖 Comprehensive documentation and quick start guide Created by FloatFloat-chan with ❤️ for the ComfyUI community
90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
import logging
|
|
import functools
|
|
import asyncio
|
|
from typing import Callable, Any
|
|
from enum import Enum
|
|
|
|
class ErrorSeverity(Enum):
|
|
LOW = "low"
|
|
MEDIUM = "medium"
|
|
HIGH = "high"
|
|
CRITICAL = "critical"
|
|
|
|
class GoogleAIPluginError(Exception):
|
|
"""Base exception for plugin errors"""
|
|
|
|
def __init__(self, message: str, severity: ErrorSeverity = ErrorSeverity.MEDIUM):
|
|
self.message = message
|
|
self.severity = severity
|
|
super().__init__(message)
|
|
|
|
class APIError(GoogleAIPluginError):
|
|
"""API-specific errors"""
|
|
pass
|
|
|
|
class RateLimitError(APIError):
|
|
"""Rate limiting errors"""
|
|
|
|
def __init__(self, message: str, retry_after: int = 60):
|
|
self.retry_after = retry_after
|
|
super().__init__(message, ErrorSeverity.HIGH)
|
|
|
|
class ConfigurationError(GoogleAIPluginError):
|
|
"""Configuration-related errors"""
|
|
pass
|
|
|
|
def error_handler(severity: ErrorSeverity = ErrorSeverity.MEDIUM,
|
|
fallback_value: Any = None):
|
|
"""Decorator for handling errors with appropriate fallbacks"""
|
|
|
|
def decorator(func: Callable):
|
|
@functools.wraps(func)
|
|
async def async_wrapper(*args, **kwargs):
|
|
try:
|
|
return await func(*args, **kwargs)
|
|
except Exception as e:
|
|
logger = logging.getLogger(func.__name__)
|
|
|
|
if isinstance(e, GoogleAIPluginError):
|
|
error_msg = f"{func.__name__}: {e.message}"
|
|
logger.log(_severity_to_log_level(e.severity), error_msg)
|
|
else:
|
|
error_msg = f"{func.__name__}: Unexpected error - {str(e)}"
|
|
logger.error(error_msg)
|
|
|
|
# Re-raise critical errors
|
|
if isinstance(e, GoogleAIPluginError) and e.severity == ErrorSeverity.CRITICAL:
|
|
raise
|
|
|
|
# Return fallback for non-critical errors
|
|
if fallback_value is not None:
|
|
return fallback_value
|
|
|
|
# Re-raise if no fallback specified
|
|
raise
|
|
|
|
@functools.wraps(func)
|
|
def sync_wrapper(*args, **kwargs):
|
|
try:
|
|
return func(*args, **kwargs)
|
|
except Exception as e:
|
|
logger = logging.getLogger(func.__name__)
|
|
logger.error(f"{func.__name__}: {str(e)}")
|
|
|
|
if fallback_value is not None:
|
|
return fallback_value
|
|
raise
|
|
|
|
return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
|
|
|
|
return decorator
|
|
|
|
def _severity_to_log_level(severity: ErrorSeverity) -> int:
|
|
"""Convert error severity to logging level"""
|
|
mapping = {
|
|
ErrorSeverity.LOW: logging.INFO,
|
|
ErrorSeverity.MEDIUM: logging.WARNING,
|
|
ErrorSeverity.HIGH: logging.ERROR,
|
|
ErrorSeverity.CRITICAL: logging.CRITICAL
|
|
}
|
|
return mapping.get(severity, logging.WARNING) |