| 1 | import logging |
| 2 | import queue |
| 3 | import sys |
| 4 | from logging.handlers import QueueHandler, QueueListener |
| 5 | from config import Config |
| 6 | |
| 7 | _listener = None |
| 8 | |
| 9 | |
| 10 | def apply_access_log_setting(): |
| 11 | """Apply the runtime access-log switch for uvicorn request logs.""" |
| 12 | logging.getLogger("uvicorn.access").disabled = not Config.ACCESS_LOG |
| 13 | |
| 14 | |
| 15 | def _managed_log_level() -> int: |
| 16 | level_name = str(getattr(Config, "LOG_LEVEL", "INFO") or "INFO").upper() |
| 17 | return getattr(logging, level_name, logging.INFO) |
| 18 | |
| 19 | |
| 20 | def apply_log_level_setting(): |
| 21 | """Apply the runtime log-level setting to the queue logger and known noisy loggers.""" |
| 22 | level = _managed_log_level() |
| 23 | root_logger = logging.getLogger() |
| 24 | root_logger.setLevel(level) |
| 25 | for handler in root_logger.handlers: |
| 26 | handler.setLevel(level) |
| 27 | if _listener is not None: |
| 28 | for handler in getattr(_listener, "handlers", []): |
| 29 | handler.setLevel(level) |
| 30 | for name in ("uvicorn", "uvicorn.error", "uvicorn.access", "httpx", "httpcore"): |
| 31 | logger = logging.getLogger(name) |
| 32 | logger.setLevel(logging.WARNING if name in {"httpx", "httpcore"} else level) |
| 33 | |
| 34 | |
| 35 | class AIGCFormatter(logging.Formatter): |
| 36 | """Custom log formatter that adds level icons for better readability.""" |
| 37 | |
| 38 | LEVEL_ICONS = { |
| 39 | "DEBUG": ".", |
| 40 | "INFO": "i", |
| 41 | "WARNING": "!", |
| 42 | "ERROR": "x", |
| 43 | "CRITICAL": "X", |
| 44 | } |
| 45 | |
| 46 | LEVEL_COLORS = { |
| 47 | "DEBUG": "\033[90m", |
| 48 | "INFO": "\033[36m", |
| 49 | "WARNING": "\033[33m", |
| 50 | "ERROR": "\033[31m", |
| 51 | "CRITICAL": "\033[1;37;41m", |
| 52 | } |
| 53 | RESET = "\033[0m" |
| 54 | |
| 55 | def __init__(self, *args, use_color=False, **kwargs): |
| 56 | super().__init__(*args, **kwargs) |
| 57 | self.use_color = use_color |
| 58 | |
| 59 | def format(self, record: logging.LogRecord) -> str: |
| 60 | record.level_icon = self.LEVEL_ICONS.get(record.levelname, record.levelname[:1]) |
| 61 | if self.use_color: |
| 62 | color = self.LEVEL_COLORS.get(record.levelname) |
| 63 | if color: |
| 64 | record.level_icon = f"{color}{record.level_icon}{self.RESET}" |
| 65 | record.levelname = f"{color}{record.levelname}{self.RESET}" |
| 66 | return super().format(record) |
| 67 | |
| 68 | |
| 69 | def setup_concurrent_logging(): |
| 70 | """Configure queue-based logging so worker threads do not interleave output.""" |
| 71 | global _listener |
| 72 | if _listener is not None: |
| 73 | return _listener |
| 74 | |
| 75 | level = _managed_log_level() |
| 76 | |
| 77 | log_queue = queue.Queue(-1) |
| 78 | console_handler = logging.StreamHandler(sys.stdout) |
| 79 | console_handler.setLevel(level) |
| 80 | console_handler.setFormatter(AIGCFormatter( |
| 81 | "%(asctime)s | %(level_icon)s %(levelname)-7s | %(name)s:%(lineno)d | %(message)s", |
| 82 | datefmt="%H:%M:%S", |
| 83 | use_color=sys.stdout.isatty(), |
| 84 | )) |
| 85 | |
| 86 | listener = QueueListener(log_queue, console_handler, respect_handler_level=True) |
| 87 | listener.start() |
| 88 | |
| 89 | queue_handler = QueueHandler(log_queue) |
| 90 | root_logger = logging.getLogger() |
| 91 | root_logger.setLevel(level) |
| 92 | root_logger.handlers.clear() |
| 93 | root_logger.addHandler(queue_handler) |
| 94 | |
| 95 | for name in ("uvicorn", "uvicorn.error", "uvicorn.access", "httpx", "httpcore"): |
| 96 | logger = logging.getLogger(name) |
| 97 | logger.handlers.clear() |
| 98 | logger.propagate = True |
| 99 | logger.setLevel(logging.WARNING if name in {"httpx", "httpcore"} else level) |
| 100 | apply_log_level_setting() |
| 101 | apply_access_log_setting() |
| 102 | |
| 103 | _listener = listener |
| 104 | return listener |
| 105 |