utils.py

"""
Utility functions for the multi-agent system including:
1. Logging setup
2. Error handling
3. Data processing helpers
4. API response validation
"""

import os
import logging
import json
import time
import requests
from typing import Any, Dict, Callable, Optional, List
from datetime import datetime
from functools import wraps
from dotenv import load_dotenv

# Custom Exceptions
class APIKeyError(Exception):
    """Raised when API key is missing or invalid"""
    pass

class APIError(Exception):
    """Raised when API call fails"""
    pass

class DataValidationError(Exception):
    """Raised when data validation fails"""
    pass

class RateLimitError(Exception):
    """Raised when API rate limit is hit"""
    pass

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def load_api_key(service_name: str) -> str:
    """
    Load API key for specified service from .env file
    Args:
        service_name: Name of the service (e.g., 'OPENWEATHERMAP', 'NEWSAPI')
    Returns:
        API key string
    Raises:
        APIKeyError: If API key not found
    """
    load_dotenv()
    key_name = f"{service_name}_API_KEY"
    api_key = os.getenv(key_name)
    
    if not api_key:
        error_msg = f"{key_name} not found in environment variables"
        logger.error(error_msg)
        raise APIKeyError(error_msg)
    
    logger.info(f"Successfully loaded API key for {service_name}")
    return api_key

def safe_api_call(max_retries: int = 3, delay: float = 1.0):
    """
    Decorator for making safe API calls with retries
    Args:
        max_retries: Maximum number of retry attempts
        delay: Delay between retries in seconds
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            last_error = None
            
            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except RateLimitError as e:
                    last_error = e
                    if attempt < max_retries - 1:
                        retry_delay = delay * (attempt + 1)  # Exponential backoff
                        logger.warning(
                            f"Rate limit hit (attempt {attempt + 1}/{max_retries}): {str(e)}. "
                            f"Retrying in {retry_delay} seconds..."
                        )
                        time.sleep(retry_delay)
                    else:
                        logger.error(f"Rate limit hit after {max_retries} attempts: {str(e)}")
                        raise last_error
                except APIError as e:
                    last_error = e
                    if attempt < max_retries - 1:
                        logger.warning(
                            f"API call failed (attempt {attempt + 1}/{max_retries}): {str(e)}. "
                            f"Retrying in {delay} seconds..."
                        )
                        time.sleep(delay)
                    else:
                        logger.error(f"API call failed after {max_retries} attempts: {str(e)}")
                        raise last_error
                except Exception as e:
                    logger.error(f"Unexpected error: {str(e)}")
                    raise
            
            return None
        return wrapper
    return decorator

def validate_required_fields(data: Dict, required_fields: List[str], context: str = "") -> None:
    """
    Validate that all required fields are present in the data
    Args:
        data: Dictionary to validate
        required_fields: List of required field names
        context: Context for error message
    Raises:
        DataValidationError: If any required field is missing
    """
    missing_fields = [field for field in required_fields if field not in data]
    if missing_fields:
        error_msg = f"Missing required fields in {context}: {missing_fields}"
        logger.error(error_msg)
        raise DataValidationError(error_msg)

def handle_api_error(response: requests.Response, context: str = "") -> None:
    """
    Handle API response errors
    Args:
        response: API response object
        context: Context for error message
    Raises:
        RateLimitError: If rate limit is hit
        APIError: For other API errors
    """
    if response.status_code == 429:
        raise RateLimitError(f"Rate limit exceeded for {context}")
    elif response.status_code >= 400:
        error_msg = f"API error for {context}: {response.status_code} - {response.text}"
        logger.error(error_msg)
        raise APIError(error_msg)

def format_timestamp(timestamp: datetime, format_str: str = "%Y-%m-%dT%H:%M:%SZ") -> str:
    """
    Format datetime object to string
    Args:
        timestamp: Datetime object to format
        format_str: Format string (default ISO 8601)
    Returns:
        Formatted timestamp string
    """
    return timestamp.strftime(format_str)

def parse_timestamp(timestamp_str: str, format_str: str = "%Y-%m-%dT%H:%M:%SZ") -> datetime:
    """
    Parse timestamp string to datetime object
    Args:
        timestamp_str: Timestamp string to parse
        format_str: Format string (default ISO 8601)
    Returns:
        Datetime object
    """
    return datetime.strptime(timestamp_str, format_str)

def validate_api_response(response: Dict[str, Any], required_fields: Optional[List[str]] = None) -> bool:
    """
    Validate API response structure and content
    Args:
        response: API response dictionary
        required_fields: List of required field names
    Returns:
        True if valid, False otherwise
    """
    if not isinstance(response, dict):
        logger.error("API response is not a dictionary")
        return False
    
    if required_fields:
        try:
            validate_required_fields(response, required_fields, "API response")
        except DataValidationError:
            return False
    
    return True

def format_coordinates(lat: float, lon: float) -> str:
    """
    Format coordinates for API requests
    Args:
        lat: Latitude
        lon: Longitude
    Returns:
        Formatted coordinate string
    """
    return f"{lat:.6f},{lon:.6f}"

def log_api_call(service: str, endpoint: str, params: Dict = None):
    """
    Log API call details
    Args:
        service: API service name
        endpoint: API endpoint
        params: Request parameters (excluding sensitive data)
    """
    log_data = {
        "service": service,
        "endpoint": endpoint,
        "timestamp": format_timestamp(datetime.now()),
        "params": params
    }
    logger.info(f"API Call: {json.dumps(log_data, indent=2)}")

def parse_datetime(date_str: str) -> datetime:
    """
    Parse datetime strings from various formats
    """
    # TODO: Implement datetime parsing
    pass