vent_agent_2.0/obs/log_setup.py

"""统一业务与服务器日志：文本便于阅读，JSON 便于采集。"""
from __future__ import annotations

import json
import logging
import re
import sys
from datetime import datetime

from obs.tracing import get_trace_id, get_log_context

_NOISY_LOGGERS = ('httpx', 'httpcore', 'urllib3', 'openai', 'aiosqlite',
                  'apscheduler.executors.default', 'apscheduler.scheduler',
                  'langchain_openai.chat_models._client_utils')
_RESERVED = {'ts', 'level', 'logger', 'trace_id', 'message', 'exc'}
_SECRET_KEYS = {'authorization', 'proxy_authorization', 'api_key', 'openai_api_key',
                'password', 'secret', 'token', 'access_token', 'refresh_token', 'cookie', 'set_cookie'}
_ASSIGNMENT = re.compile(r'(?i)((?:api[_-]?key|password|access[_-]?token|refresh[_-]?token|secret)\s*[=:]\s*[\"\']?)[^\s,;\"\'&]+')
_BEARER = re.compile(r'(?i)(bearer\s+)[A-Za-z0-9._~+/-]+=*')
_URL_CREDENTIALS = re.compile(r'(://)[^\s/@]+:[^\s/@]+@')


def _safe(value):
    if isinstance(value, dict):
        return {str(k): '[REDACTED]' if str(k).lower().replace('-', '_') in _SECRET_KEYS else _safe(v)
                for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_safe(item) for item in value]
    if isinstance(value, str):
        value = _ASSIGNMENT.sub(r'\1[REDACTED]', value)
        value = _BEARER.sub(r'\1[REDACTED]', value)
        return _URL_CREDENTIALS.sub(r'\1[REDACTED]@', value)
    return value


class _TraceIdFilter(logging.Filter):
    def filter(self, record):
        record.trace_id = get_trace_id() or '-'
        record.log_context = get_log_context()
        return True


def _payload(record):
    payload = {
        'ts': datetime.fromtimestamp(record.created).astimezone().isoformat(timespec='milliseconds'),
        'level': record.levelname,
        'logger': record.name,
        'trace_id': getattr(record, 'trace_id', '-'),
        'message': _safe(record.getMessage()),
    }
    context = getattr(record, 'log_context', {})
    kv = getattr(record, 'kv', {})
    if isinstance(kv, dict):
        payload.update({k: v for k, v in _safe(kv).items() if k not in _RESERVED})
    # 绑定的执行身份优先于临时字段，防止覆盖关联信息。
    payload.update({k: v for k, v in _safe(context).items() if k not in _RESERVED})
    if record.exc_info:
        payload['exc'] = _safe(logging.Formatter().formatException(record.exc_info))
    if record.stack_info:
        payload['stack'] = _safe(record.stack_info)
    return payload


class _JsonFormatter(logging.Formatter):
    def format(self, record):
        return json.dumps(_payload(record), ensure_ascii=False, default=str)


class _TextFormatter(logging.Formatter):
    def format(self, record):
        p = _payload(record)
        # JSON 转义保持一条普通日志一行；异常堆栈在下一行完整呈现。
        message = str(p['message']).replace('\r', r'\r').replace('\n', r'\n')
        fields = ' '.join(f'{key}={json.dumps(value, ensure_ascii=False, default=str)}'
                          for key, value in p.items() if key not in _RESERVED)
        line = f"{p['ts']} {p['level']:<7} [{p['trace_id']}] {p['logger']} | {message}"
        if fields:
            line += ' | ' + fields
        if p.get('exc'):
            line += '\n' + p['exc']
        return line


def setup_logging(level: str = 'INFO', fmt: str = 'text') -> None:
    """重复配置只替换自身 handler，保留测试/宿主添加的采集器。"""
    root = logging.getLogger()
    root.setLevel(getattr(logging, str(level).upper(), logging.INFO))
    for handler in list(root.handlers):
        if getattr(handler, '_vent_handler', False):
            root.removeHandler(handler)
            handler.close()
    handler = logging.StreamHandler(sys.stdout)
    handler._vent_handler = True
    handler.addFilter(_TraceIdFilter())
    handler.setFormatter(_JsonFormatter() if str(fmt).lower() == 'json' else _TextFormatter())
    root.addHandler(handler)
    for name in _NOISY_LOGGERS:
        logging.getLogger(name).setLevel(logging.WARNING)
    for name in ('uvicorn', 'uvicorn.error', 'uvicorn.access'):
        server = logging.getLogger(name)
        server.handlers.clear()
        server.propagate = True
        server.setLevel(logging.NOTSET)
    # 请求汇总由 ASGI 中间件记录，包含流结束时间，避免 Uvicorn 再记一遍。
    logging.getLogger('uvicorn.access').disabled = True


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