Files
tw/core/websocket_logger_setup.py

136 lines
4.3 KiB
Python

import logging
import os
import subprocess
from datetime import datetime
from pathlib import Path
class _AnsiColorFormatter(logging.Formatter):
RESET = "\033[0m"
MESSAGE_TEXT_REPLACEMENTS = (
("[PROMPT->AI 前置提示词]", "[AI提示词]"),
("[PROMPT->AI", "[AI提示词"),
("[THINK/TOOL_CALL]", "[AI思考-工具调用]"),
("[THINK/TOOL_RETURN]", "[AI思考-工具返回]"),
("[THINK/RAW_OUTPUT]", "[AI思考-原始输出]"),
("[REPLY->CUSTOMER]", "[AI回复客户]"),
("[ACTIVITY]", "[活动日志]"),
("[AI质检]", "[AI质检]"),
)
MESSAGE_COLOR_RULES = (
("[PROMPT->AI", "\033[94m"),
("[THINK/", "\033[96m"),
("[REPLY->CUSTOMER]", "\033[92m"),
("Agent 回复", "\033[92m"),
("[ACTIVITY]", "\033[95m"),
("[AI质检]", "\033[97m"),
("收到新消息", "\033[36m"),
("发送成功", "\033[32m"),
("防抖等待", "\033[93m"),
)
COLORS = {
logging.DEBUG: "\033[36m",
logging.INFO: "\033[32m",
logging.WARNING: "\033[33m",
logging.ERROR: "\033[31m",
logging.CRITICAL: "\033[35m",
}
def __init__(self, fmt: str, datefmt: str | None = None, use_color: bool = True):
super().__init__(fmt=fmt, datefmt=datefmt)
self.use_color = use_color
def format(self, record: logging.LogRecord) -> str:
msg = super().format(record)
if not self.use_color:
for old, new in self.MESSAGE_TEXT_REPLACEMENTS:
msg = msg.replace(old, new)
return msg
raw_msg = record.getMessage()
for old, new in self.MESSAGE_TEXT_REPLACEMENTS:
msg = msg.replace(old, new)
for key, color in self.MESSAGE_COLOR_RULES:
if key in raw_msg:
return f"{color}{msg}{self.RESET}"
color = self.COLORS.get(record.levelno, "")
if not color:
return msg
return f"{color}{msg}{self.RESET}"
_APP_VERSION = None
_LOG_RECORD_FACTORY_INSTALLED = False
def get_app_log_version() -> str:
global _APP_VERSION
if _APP_VERSION:
return _APP_VERSION
env_version = str(os.getenv("APP_VERSION", "")).strip()
if env_version:
_APP_VERSION = env_version
return _APP_VERSION
try:
repo_root = Path(__file__).resolve().parent.parent
git_version = subprocess.check_output(
["git", "-C", str(repo_root), "rev-parse", "--short", "HEAD"],
stderr=subprocess.DEVNULL,
text=True,
).strip()
except Exception:
git_version = ""
_APP_VERSION = git_version or "dev"
os.environ.setdefault("APP_VERSION", _APP_VERSION)
return _APP_VERSION
def install_log_record_factory():
global _LOG_RECORD_FACTORY_INSTALLED
if _LOG_RECORD_FACTORY_INSTALLED:
return
version = get_app_log_version()
old_factory = logging.getLogRecordFactory()
def record_factory(*args, **kwargs):
record = old_factory(*args, **kwargs)
record.app_version = getattr(record, "app_version", version)
return record
logging.setLogRecordFactory(record_factory)
_LOG_RECORD_FACTORY_INSTALLED = True
def setup_logger():
from logging.handlers import RotatingFileHandler
from config.config import LOG_DIR, LOG_MAX_BYTES, LOG_BACKUP_COUNT
install_log_record_factory()
logger = logging.getLogger("cs_agent")
if getattr(logger, "_cs_logger_configured", False):
return logger
logger.setLevel(logging.INFO)
logger.propagate = False
fmt = logging.Formatter("[v%(app_version)s][%(asctime)s] %(message)s", datefmt="%H:%M:%S")
use_color = (os.getenv("LOG_COLOR", "1").lower() in ("1", "true", "yes")) and not bool(os.getenv("NO_COLOR"))
ch = logging.StreamHandler()
ch.setFormatter(_AnsiColorFormatter("[v%(app_version)s][%(asctime)s] %(message)s", datefmt="%H:%M:%S", use_color=use_color))
logger.addHandler(ch)
LOG_DIR.mkdir(exist_ok=True)
today = datetime.now().strftime("%Y-%m-%d")
fh = RotatingFileHandler(
LOG_DIR / f"chat_{today}.log",
maxBytes=LOG_MAX_BYTES,
backupCount=LOG_BACKUP_COUNT,
encoding="utf-8",
)
fh.setFormatter(fmt)
logger.addHandler(fh)
logger._cs_logger_configured = True
return logger