| import os |
| import logging |
| import colorama |
| from logging.handlers import RotatingFileHandler |
| from typing_extensions import Literal |
|
|
| |
| colorama.init() |
|
|
| log_level = Literal[ |
| "CRITICAL", |
| "ERROR", |
| "WARNING", |
| "INFO", |
| "DEBUG", |
| ] |
|
|
|
|
| def get_logging_level(level: str) -> int: |
| levels = { |
| "CRITICAL": logging.CRITICAL, |
| "ERROR": logging.ERROR, |
| "WARNING": logging.WARNING, |
| "INFO": logging.INFO, |
| "DEBUG": logging.DEBUG, |
| } |
| return levels.get(level.upper(), logging.INFO) |
|
|
|
|
| class ColorFormatter(logging.Formatter): |
| LEVEL_COLORS = { |
| logging.DEBUG: colorama.Fore.BLUE, |
| logging.INFO: colorama.Fore.GREEN, |
| logging.WARNING: colorama.Fore.YELLOW, |
| logging.ERROR: colorama.Fore.RED, |
| logging.CRITICAL: colorama.Fore.LIGHTRED_EX, |
| } |
| RESET = colorama.Style.RESET_ALL |
|
|
| def format(self, record): |
| color = self.LEVEL_COLORS.get(record.levelno, self.RESET) |
| record.levelname = f"{color}{record.levelname}{self.RESET}" |
| return super().format(record) |
|
|
|
|
| def setup_logger( |
| log_level: log_level | str = "INFO", save_logs: bool = False |
| ) -> logging.Logger: |
| logger = logging.getLogger("qwen_api") |
| logger.propagate = False |
|
|
| if logger.handlers: |
| logger.handlers.clear() |
|
|
| logger.setLevel(get_logging_level(log_level)) |
|
|
| base_format = "[%(levelname)s] %(asctime)s - %(name)s -> %(message)s" |
|
|
| |
| console_handler = logging.StreamHandler() |
| console_handler.setFormatter(ColorFormatter(base_format)) |
| logger.addHandler(console_handler) |
|
|
| |
| if save_logs: |
| log_dir = "logs" |
| os.makedirs(log_dir, exist_ok=True) |
| file_handler = RotatingFileHandler( |
| f"{log_dir}/qwen.log", maxBytes=5 * 1024 * 1024, backupCount=3 |
| ) |
| file_handler.setFormatter(logging.Formatter(base_format)) |
| logger.addHandler(file_handler) |
|
|
| return logger |
|
|