agent_jrxml/backend/logger.py

"""集中日志模块。

提供：
- 结构化 JSON 日志（每行一条记录）
- 请求级 trace_id（通过 contextvars 自动传播）
- 独立的 LLM 调用日志文件
- 日志轮转（按大小 10MB，保留 5 个备份）

用法:
    from backend.logger import get_logger, set_trace_id

    # 业务日志
    log = get_logger("agent")
    log.info("节点开始执行", extra={"node": "classify_intent", "session_id": "xxx"})

    # LLM 日志
    llm_log = get_logger("llm")
    llm_log.info("LLM 请求", extra={"prompt": "...", "model": "gpt-4o"})
"""

import json
import logging
import os
import sys
import uuid
from contextvars import ContextVar
from datetime import datetime, timezone, timedelta
from logging.handlers import RotatingFileHandler
from pathlib import Path
from typing import Optional

LOG_DIR = Path(os.getenv("LOG_DIR", "./logs"))
LOG_LEVEL = os.getenv("LOG_LEVEL", "DEBUG")
LLM_LOG_FILE = "llm.log"
APP_LOG_FILE = "app.log"

CHINA_TZ = timezone(timedelta(hours=8))

_trace_id_var: ContextVar[str] = ContextVar("trace_id", default="")


def generate_trace_id() -> str:
    return uuid.uuid4().hex[:16]


def get_trace_id() -> str:
    tid = _trace_id_var.get()
    if not tid:
        tid = generate_trace_id()
        _trace_id_var.set(tid)
    return tid


def set_trace_id(trace_id: str):
    _trace_id_var.set(trace_id)


class JsonFormatter(logging.Formatter):
    """将日志记录格式化为单行 JSON，便于后续分析。

    LogRecord 标准属性的键（不放入 extra）。
    通过 logging.Logger.debug(msg, extra={...}) 传入的键会自动设为
    LogRecord 属性，由本格式化器收集到 extra 字段中。
    """

    _STANDARD_ATTRS: set[str] = frozenset({
        "args", "asctime", "created", "exc_info", "exc_text", "filename",
        "funcName", "levelname", "levelno", "lineno", "module", "msecs",
        "message", "msg", "name", "pathname", "process", "processName",
        "relativeCreated", "stack_info", "thread", "threadName",
        "extra_fields", "taskName",
    })

    def _collect_extra(self, record: logging.LogRecord) -> dict:
        """从 LogRecord 上收集非标准属性 → 合并为 extra dict。"""
        extra = dict(getattr(record, "extra_fields", {}))
        for key, val in record.__dict__.items():
            if key not in self._STANDARD_ATTRS and not key.startswith("_"):
                extra[key] = val
        return extra

    def format(self, record: logging.LogRecord) -> str:
        log_entry = {
            "timestamp": datetime.now(CHINA_TZ).isoformat(),
            "level": record.levelname,
            "logger": record.name,
            "trace_id": get_trace_id(),
            "message": record.getMessage(),
            "module": record.module,
            "function": record.funcName,
            "line": record.lineno,
        }

        extra = self._collect_extra(record)
        if extra:
            log_entry["extra"] = extra

        if record.exc_info and record.exc_info[0]:
            import traceback
            log_entry["exception"] = traceback.format_exception(
                record.exc_info[0], record.exc_info[1], record.exc_info[2]
            )

        return json.dumps(log_entry, ensure_ascii=False)


def _create_handler(filename: str, level: int) -> RotatingFileHandler:
    handler = RotatingFileHandler(
        filename=str(LOG_DIR / filename),
        maxBytes=10 * 1024 * 1024,
        backupCount=5,
        encoding="utf-8",
    )
    handler.setLevel(level)
    handler.setFormatter(JsonFormatter())
    return handler


def _get_level() -> int:
    return getattr(logging, LOG_LEVEL.upper(), logging.DEBUG)


def get_logger(name: str) -> logging.Logger:
    """获取指定名称的 logger，自动配置了 JSON 格式化 + 文件轮转。

    name="llm" → 输出到 logs/llm.log（仅 LLM 调用相关）
    其他 name → 输出到 logs/app.log
    """
    logger = logging.getLogger(f"jrxml.{name}")

    if logger.handlers:
        return logger

    LOG_DIR.mkdir(parents=True, exist_ok=True)

    level = _get_level()
    logger.setLevel(level)
    logger.propagate = False

    if name == "llm":
        logger.addHandler(_create_handler(LLM_LOG_FILE, level))
    else:
        logger.addHandler(_create_handler(APP_LOG_FILE, level))

    return logger


class _ExtraAdapter(logging.LoggerAdapter):
    """支持通过 adapter.extra 合并 extra 字段的适配器。"""

    def process(self, msg, kwargs):
        extra = kwargs.pop("extra", {})
        merged = {**self.extra, **extra} if self.extra or extra else None
        if merged:
            kwargs["extra"] = {"extra_fields": merged}
        return msg, kwargs


def get_trace_logger(name: str) -> _ExtraAdapter:
    """返回一个自动附带 trace_id 的 logger 适配器。

    用法:
        log = get_trace_logger("agent")
        log.info("节点完成", extra={"node": "generate"})
    """
    logger = get_logger(name)
    return _ExtraAdapter(logger, {"trace_id": get_trace_id()})