AgentCoreService/pqai_agent/logging.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8

import logging
from logging.handlers import RotatingFileHandler

from pqai_agent import configs

COLORS = {
    'DEBUG': '\033[0m',    # 蓝色
    'INFO': '\033[32m',     # 绿色
    'WARNING': '\x1b[33;20m',  # 黄色
    'ERROR': '\033[91m',    # 红色
    'CRITICAL': '\033[95m', # 紫色
    'RESET': '\033[0m',     # 重置颜色
}

class ColoredFormatter(logging.Formatter):
    def format(self, record):
        # 获取原始日志消息
        message = super().format(record)
        # 根据日志等级添加颜色
        if record.levelname in COLORS:
            message = f"{COLORS[record.levelname]}{message}{COLORS['RESET']}"
        return message

def setup_root_logger(level=logging.DEBUG, logfile_name='service.log'):
    logging_format = '%(asctime)s - %(name)s %(funcName)s[%(lineno)d] - %(levelname)s - %(message)s'
    plain_formatter = logging.Formatter(logging_format)
    color_formatter = ColoredFormatter(logging_format)
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.DEBUG)
    console_handler.setFormatter(color_formatter)

    root_logger = logging.getLogger()
    root_logger.handlers.clear()
    root_logger.addHandler(console_handler)

    log_dir = configs.get().get('system', {}).get('log_dir', '')
    if log_dir:
        file_handler = RotatingFileHandler(
            f'{log_dir}/{logfile_name}',
            maxBytes=64 * 1024 * 1024,
            backupCount=5,
            encoding='utf-8'
        )
        file_handler.setLevel(logging.DEBUG)
        file_handler.setFormatter(plain_formatter)
        root_logger.addHandler(file_handler)

    agent_logger = logging.getLogger('agent')
    agent_logger.setLevel(level)

logger = logging.getLogger('agent')