salim0986/graph-bug-ai

1import logging2import sys3import os4import json5from datetime import datetime6from typing import Any, Dict, Optional7from .config import LOG_LEVEL8 9class StructuredFormatter(logging.Formatter):10    """11    JSON formatter for structured logging12    Enables better log aggregation and analysis in production13    """14    15    def format(self, record: logging.LogRecord) -> str:16        """Format log record as JSON"""17        log_data: Dict[str, Any] = {18            "timestamp": datetime.utcnow().isoformat() + "Z",19            "level": record.levelname,20            "logger": record.name,21            "message": record.getMessage(),22            "module": record.module,23            "function": record.funcName,24            "line": record.lineno,25        }26        27        # Add exception info if present28        if record.exc_info:29            log_data["exception"] = self.formatException(record.exc_info)30        31        # Add any custom fields from extra parameter32        if hasattr(record, "extra_fields"):33            log_data.update(record.extra_fields)34        35        # Add context fields if present36        for key in ["repo_id", "installation_id", "user_id", "request_id", "endpoint"]:37            if hasattr(record, key):38                log_data[key] = getattr(record, key)39        40        return json.dumps(log_data)41 42 43def setup_logger(name: str, structured: bool = False) -> logging.Logger:44    """45    Setup a logger with consistent formatting.46    47    Args:48        name: Logger name (usually __name__)49        structured: Use JSON structured logging (recommended for production)50    51    Returns:52        Configured logger instance53    """54    logger = logging.getLogger(name)55    56    # Only configure if not already configured57    if not logger.handlers:58        logger.setLevel(getattr(logging, LOG_LEVEL))59        60        # Detect production environment (Hugging Face Spaces, Cloud Run, etc.)61        is_production = os.getenv("ENVIRONMENT") == "production" or os.getenv("SPACE_ID") is not None62        63        # Console handler with immediate flushing for production64        handler = logging.StreamHandler(sys.stdout)65        handler.setLevel(getattr(logging, LOG_LEVEL))66        67        # Force immediate flush in production to ensure logs are visible68        if is_production:69            # Set unbuffered mode70            sys.stdout.reconfigure(line_buffering=True) if hasattr(sys.stdout, 'reconfigure') else None71        72        # Choose formatter based on environment73        if structured or is_production:74            # Use structured JSON logging in production75            formatter = StructuredFormatter()76        else:77            # Use human-readable format for development78            formatter = logging.Formatter(79                '%(asctime)s - %(name)s - %(levelname)s - %(message)s',80                datefmt='%Y-%m-%d %H:%M:%S'81            )82        83        handler.setFormatter(formatter)84        logger.addHandler(handler)85        86        # Force flush after each log in production87        if is_production:88            old_emit = handler.emit89            def flush_emit(record):90                old_emit(record)91                handler.flush()92            handler.emit = flush_emit93    94    return logger95 96 97class LogContext:98    """99    Context manager for adding structured fields to all logs within a block100    101    Usage:102        with LogContext(logger, repo_id="123", installation_id="456"):103            logger.info("Processing repo")  # Will include repo_id and installation_id104    """105    106    def __init__(self, logger: logging.Logger, **kwargs):107        self.logger = logger108        self.context = kwargs109        self.old_factory = None110    111    def __enter__(self):112        old_factory = logging.getLogRecordFactory()113        114        def record_factory(*args, **kwargs):115            record = old_factory(*args, **kwargs)116            for key, value in self.context.items():117                setattr(record, key, value)118            return record119        120        logging.setLogRecordFactory(record_factory)121        self.old_factory = old_factory122        return self123    124    def __exit__(self, exc_type, exc_val, exc_tb):125        if self.old_factory:126            logging.setLogRecordFactory(self.old_factory)127