report_generation/log/logger.py

from __future__ import annotations

import logging
from logging.handlers import RotatingFileHandler
from pathlib import Path


_CONFIGURED = False
_FILE_PROCESSING_PREFIXES = (
    "worker.document_processing",
    "services.kb_service",
    "services.es_docs",
    "services.element_llm_extract_service",
    "routers.extract",
    "function.documents",
    "function.vector_store",
    "repo.kb_documents",
    "routers.reference",
    "services.doc_convert_service",
    "services.reference_service",
)
_DOCUMENT_GENERATION_PREFIXES = (
    "services.write_service",
    "services.report_generation_service",
    "services.markdown_stream_service",
    "services.llm_client",
    "services.llm_runner",
    "services.report_prompt_service",
    "services.report_runtime_store",
)
# 生成全过程追踪：完整记录输入 prompt / 调用模型 / 模型输出
_GENERATION_TRACE_PREFIXES = (
    "generation.trace",
)


class _PrefixFilter(logging.Filter):
    def __init__(self, prefixes: tuple[str, ...]) -> None:
        super().__init__()
        self.prefixes = prefixes

    def filter(self, record: logging.LogRecord) -> bool:
        name = str(record.name or "")
        return any(name == prefix or name.startswith(prefix + ".") for prefix in self.prefixes)


class _OtherFilter(logging.Filter):
    def filter(self, record: logging.LogRecord) -> bool:
        name = str(record.name or "")
        if any(name == prefix or name.startswith(prefix + ".") for prefix in _FILE_PROCESSING_PREFIXES):
            return False
        if any(name == prefix or name.startswith(prefix + ".") for prefix in _DOCUMENT_GENERATION_PREFIXES):
            return False
        if any(name == prefix or name.startswith(prefix + ".") for prefix in _GENERATION_TRACE_PREFIXES):
            return False
        return True


def configure_logging(
    *,
    log_dir: str | Path = "logs",
    level: int = logging.INFO,
) -> Path:
    global _CONFIGURED

    target_dir = Path(log_dir).resolve()
    target_dir.mkdir(parents=True, exist_ok=True)
    other_log_path = target_dir / "other.log"

    if _CONFIGURED:
        return other_log_path

    formatter = logging.Formatter(
        "%(asctime)s | %(levelname)s | %(name)s | %(message)s"
    )

    root_logger = logging.getLogger()
    root_logger.setLevel(level)

    file_processing_handler = RotatingFileHandler(
        target_dir / "file_processing.log",
        maxBytes=10 * 1024 * 1024,
        backupCount=5,
        encoding="utf-8",
    )
    file_processing_handler.setLevel(level)
    file_processing_handler.setFormatter(formatter)
    file_processing_handler.addFilter(_PrefixFilter(_FILE_PROCESSING_PREFIXES))

    document_generation_handler = RotatingFileHandler(
        target_dir / "document_generation.log",
        maxBytes=10 * 1024 * 1024,
        backupCount=5,
        encoding="utf-8",
    )
    document_generation_handler.setLevel(level)
    document_generation_handler.setFormatter(formatter)
    document_generation_handler.addFilter(_PrefixFilter(_DOCUMENT_GENERATION_PREFIXES))

    other_handler = RotatingFileHandler(
        other_log_path,
        maxBytes=10 * 1024 * 1024,
        backupCount=5,
        encoding="utf-8",
    )
    other_handler.setLevel(level)
    other_handler.setFormatter(formatter)
    other_handler.addFilter(_OtherFilter())

    # ── 要素抽取独立日志 ─────────────────────────────────────────────
    element_extract_handler = RotatingFileHandler(
        target_dir / "element_extract.log",
        maxBytes=10 * 1024 * 1024,
        backupCount=10,
        encoding="utf-8",
    )
    element_extract_handler.setLevel(level)
    element_extract_handler.setFormatter(formatter)
    element_extract_handler.addFilter(_PrefixFilter(("services.element_llm_extract_service", "routers.extract")))

    # ── 文件上传/解析独立日志 ─────────────────────────────────────────
    file_upload_handler = RotatingFileHandler(
        target_dir / "file_upload.log",
        maxBytes=10 * 1024 * 1024,
        backupCount=10,
        encoding="utf-8",
    )
    file_upload_handler.setLevel(level)
    file_upload_handler.setFormatter(formatter)
    file_upload_handler.addFilter(_PrefixFilter(("routers.reference", "routers.template", "services.doc_convert_service", "services.reference_service", "services.kb_service", "routers.kb")))

    # ── 报告生成独立日志 ──────────────────────────────────────────────
    report_generation_handler = RotatingFileHandler(
        target_dir / "report_generation.log",
        maxBytes=10 * 1024 * 1024,
        backupCount=10,
        encoding="utf-8",
    )
    report_generation_handler.setLevel(level)
    report_generation_handler.setFormatter(formatter)
    report_generation_handler.addFilter(_PrefixFilter(("services.report_generation_service", "services.report_prompt_service", "services.report_runtime_store", "services.markdown_stream_service")))

    # ── LLM 调用独立日志 ──────────────────────────────────────────────
    llm_handler = RotatingFileHandler(
        target_dir / "llm.log",
        maxBytes=10 * 1024 * 1024,
        backupCount=10,
        encoding="utf-8",
    )
    llm_handler.setLevel(level)
    llm_handler.setFormatter(formatter)
    llm_handler.addFilter(_PrefixFilter(("services.llm_client", "services.llm_runner")))

    # ── 生成全过程追踪日志（输入 prompt / 模型 / 输出，单条可能较大）────────
    generation_trace_handler = RotatingFileHandler(
        target_dir / "generation_trace.log",
        maxBytes=50 * 1024 * 1024,
        backupCount=10,
        encoding="utf-8",
    )
    generation_trace_handler.setLevel(level)
    generation_trace_handler.setFormatter(formatter)
    generation_trace_handler.addFilter(_PrefixFilter(_GENERATION_TRACE_PREFIXES))

    stream_handler = logging.StreamHandler()
    stream_handler.setLevel(level)
    stream_handler.setFormatter(formatter)

    root_logger.handlers.clear()
    root_logger.addHandler(file_processing_handler)
    root_logger.addHandler(document_generation_handler)
    root_logger.addHandler(other_handler)
    root_logger.addHandler(element_extract_handler)
    root_logger.addHandler(file_upload_handler)
    root_logger.addHandler(report_generation_handler)
    root_logger.addHandler(llm_handler)
    root_logger.addHandler(generation_trace_handler)
    root_logger.addHandler(stream_handler)

    _CONFIGURED = True
    return other_log_path


def get_logger(name: str) -> logging.Logger:
    return logging.getLogger(name)