LQAgentPlatform/foundation/observability/monitoring/rag/rag_monitor.py

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

"""
RAG链路监控装饰器
支持同步/异步函数的输入输出监控
"""

import time
import json
import functools
import inspect
from typing import Callable, Optional, Dict, Any, Union
from pathlib import Path

from foundation.observability.logger.loggering import review_logger as logger


class RAGMonitor:
    """RAG监控管理器"""

    def __init__(self, save_dir: str = "temp/rag_monitoring"):
        """
        初始化RAG监控器

        Args:
            save_dir: 监控数据保存目录
        """
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
        self.pipeline_data = {}
        self.current_trace_id = None

    def start_trace(self, trace_id: str, metadata: Optional[Dict] = None):
        """
        开始一个新的追踪会话

        Args:
            trace_id: 追踪会话ID
            metadata: 会话元数据
        """
        self.current_trace_id = trace_id
        self.pipeline_data[trace_id] = {
            "trace_id": trace_id,
            "start_time": time.time(),
            "metadata": metadata or {},
            "steps": {}
        }
        logger.info(f"[RAG监控] 开始追踪会话: {trace_id}")

    def end_trace(self, trace_id: str = None) -> Optional[Dict]:
        """
        结束追踪会话并保存数据

        Args:
            trace_id: 追踪会话ID，如果为None则使用当前会话

        Returns:
            追踪会话数据
        """
        trace_id = trace_id or self.current_trace_id
        if trace_id not in self.pipeline_data:
            logger.warning(f"[RAG监控] 追踪会话不存在: {trace_id}")
            return None

        data = self.pipeline_data[trace_id]
        data["end_time"] = time.time()
        data["total_duration"] = round(data["end_time"] - data["start_time"], 3)

        # 保存到文件
        file_path = self.save_dir / f"{trace_id}.json"
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2, default=str)

        logger.info(f"[RAG监控] 追踪会话已保存: {file_path}, 总耗时: {data['total_duration']}秒")
        return data

    def get_trace_data(self, trace_id: str = None) -> Optional[Dict]:
        """
        获取追踪会话数据

        Args:
            trace_id: 追踪会话ID，如果为None则使用当前会话

        Returns:
            追踪会话数据
        """
        trace_id = trace_id or self.current_trace_id
        return self.pipeline_data.get(trace_id)

    def monitor_step(
        self,
        step_name: str,
        capture_input: bool = True,
        capture_output: bool = True,
        input_transform: Optional[Callable] = None,
        output_transform: Optional[Callable] = None,
        max_input_length: int = 500,
        max_output_length: int = 1000
    ):
        """
        监控装饰器 - 支持同步和异步函数

        Args:
            step_name: 步骤名称
            capture_input: 是否捕获输入参数
            capture_output: 是否捕获输出结果
            input_transform: 输入数据转换函数（用于过滤敏感信息或压缩数据）
            output_transform: 输出数据转换函数
            max_input_length: 输入数据最大长度（超过会截断）
            max_output_length: 输出数据最大长度（超过会截断）

        Example:
            @rag_monitor.monitor_step("query_extract", capture_input=True)
            def extract_query(content: str):
                return query_rewrite_manager.query_extract(content)
        """
        def decorator(func: Callable):
            # 判断是否为异步函数
            is_async = inspect.iscoroutinefunction(func)

            if is_async:
                @functools.wraps(func)
                async def async_wrapper(*args, **kwargs):
                    return await self._execute_with_monitoring(
                        func, step_name, args, kwargs,
                        capture_input, capture_output,
                        input_transform, output_transform,
                        max_input_length, max_output_length,
                        is_async=True
                    )
                return async_wrapper
            else:
                @functools.wraps(func)
                def sync_wrapper(*args, **kwargs):
                    return self._execute_with_monitoring(
                        func, step_name, args, kwargs,
                        capture_input, capture_output,
                        input_transform, output_transform,
                        max_input_length, max_output_length,
                        is_async=False
                    )
                return sync_wrapper

        return decorator

    def _execute_with_monitoring(
        self,
        func: Callable,
        step_name: str,
        args: tuple,
        kwargs: dict,
        capture_input: bool,
        capture_output: bool,
        input_transform: Optional[Callable],
        output_transform: Optional[Callable],
        max_input_length: int,
        max_output_length: int,
        is_async: bool
    ):
        """执行函数并监控"""
        trace_id = self.current_trace_id
        if not trace_id:
            logger.warning(f"[RAG监控] 未找到活跃的追踪会话，跳过监控: {step_name}")
            # 即使没有追踪会话，也要正常执行函数
            if is_async:
                import asyncio
                return asyncio.create_task(func(*args, **kwargs))
            else:
                return func(*args, **kwargs)

        # 记录步骤数据
        step_data = {
            "step_name": step_name,
            "function_name": func.__name__,
            "start_time": time.time()
        }

        # 捕获输入
        if capture_input:
            input_data = {
                "args": self._safe_serialize(args, max_input_length),
                "kwargs": self._safe_serialize(kwargs, max_input_length)
            }
            if input_transform:
                try:
                    input_data = input_transform(input_data)
                except Exception as e:
                    logger.warning(f"[RAG监控] 输入转换失败: {e}")
            step_data["input"] = input_data

        # 执行函数
        try:
            if is_async:
                # 对于异步函数，需要特殊处理
                import asyncio

                async def async_exec():
                    result = await func(*args, **kwargs)
                    self._finalize_step(step_data, result, trace_id, capture_output,
                                       output_transform, max_output_length, success=True)
                    return result

                return asyncio.create_task(async_exec())
            else:
                result = func(*args, **kwargs)
                self._finalize_step(step_data, result, trace_id, capture_output,
                                   output_transform, max_output_length, success=True)
                return result

        except Exception as e:
            self._finalize_step(step_data, None, trace_id, capture_output,
                               output_transform, max_output_length, success=False, error=e)
            raise

    def _finalize_step(
        self,
        step_data: Dict,
        result: Any,
        trace_id: str,
        capture_output: bool,
        output_transform: Optional[Callable],
        max_output_length: int,
        success: bool,
        error: Optional[Exception] = None
    ):
        """完成步骤监控数据记录"""
        if success:
            step_data["status"] = "success"

            # 捕获输出
            if capture_output:
                output_data = self._safe_serialize(result, max_output_length)
                if output_transform:
                    try:
                        output_data = output_transform(output_data)
                    except Exception as e:
                        logger.warning(f"[RAG监控] 输出转换失败: {e}")
                step_data["output"] = output_data
        else:
            step_data["status"] = "error"
            step_data["error"] = {
                "type": type(error).__name__,
                "message": str(error)
            }
            logger.error(f"[RAG监控] 步骤执行失败: {step_data['step_name']}, 错误: {error}")

        step_data["end_time"] = time.time()
        step_data["duration"] = round(step_data["end_time"] - step_data["start_time"], 3)

        # 保存步骤数据
        if trace_id in self.pipeline_data:
            # 如果步骤名称已存在，添加序号
            original_step_name = step_data['step_name']
            step_name = original_step_name
            counter = 1
            while step_name in self.pipeline_data[trace_id]["steps"]:
                step_name = f"{original_step_name}_{counter}"
                counter += 1

            self.pipeline_data[trace_id]["steps"][step_name] = step_data
            logger.info(f"[RAG监控] 步骤完成: {step_name}, 耗时: {step_data['duration']}秒")

    def _safe_serialize(self, obj: Any, max_length: int = 500) -> Any:
        """
        安全序列化对象（防止大对象占用过多内存）

        Args:
            obj: 要序列化的对象
            max_length: 字符串最大长度

        Returns:
            序列化后的对象
        """
        if obj is None:
            return None

        # 基本类型直接返回
        if isinstance(obj, (int, float, bool)):
            return obj

        if isinstance(obj, str):
            if len(obj) > max_length:
                return {
                    "type": "string",
                    "length": len(obj),
                    "preview": obj[:max_length],
                    "truncated": True
                }
            return obj

        # 列表类型
        if isinstance(obj, (list, tuple)):
            result = {
                "type": "list" if isinstance(obj, list) else "tuple",
                "count": len(obj)
            }

            # 只保留前3项的预览
            if len(obj) > 0:
                result["preview"] = [self._safe_serialize(item, max_length) for item in obj[:3]]

            # 如果列表项少于10项，保存完整数据
            if len(obj) <= 10:
                result["full_data"] = [self._safe_serialize(item, max_length) for item in obj]
            else:
                result["truncated"] = True

            return result

        # 字典类型
        if isinstance(obj, dict):
            result = {}
            keys_list = list(obj.keys())

            # 最多保留20个键
            for key in keys_list[:20]:
                try:
                    result[str(key)] = self._safe_serialize(obj[key], max_length)
                except Exception as e:
                    result[str(key)] = f"<序列化失败: {e}>"

            if len(keys_list) > 20:
                result["__truncated__"] = f"省略了 {len(keys_list) - 20} 个键"

            return result

        # 其他类型尝试转为字符串
        try:
            str_repr = str(obj)
            if len(str_repr) > max_length:
                return {
                    "type": type(obj).__name__,
                    "preview": str_repr[:max_length],
                    "truncated": True
                }
            return {"type": type(obj).__name__, "value": str_repr}
        except Exception as e:
            return {"type": type(obj).__name__, "error": f"无法序列化: {e}"}


# 全局监控实例
rag_monitor = RAGMonitor()