第7章: 「数析」智能数据分析台

# ============================
# 目的：用“工具化”的方式把 Pandas 能力做成可控产品，而不是让 LLM 直接写/执行任意脚本
#
# 你将看到：
# - 用 @tool 把“看表结构 / 统计指标 / 画图”等能力封装成受控工具
# - ReAct Agent 通过 Thought/Action/Observation 循环选择调用这些工具
# - verbose=True 时可以观察每一步工具调用轨迹
#
# 注意：本段是“最小可运行示例”，为了教学清晰做了简化：
# - 数据用内存 DataFrame 模拟
# - 绘图输出保存到本地文件
# ============================

import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm
from langchain.tools import tool
from langchain_classic.agents import AgentExecutor, create_react_agent
from langchain_community.chat_models import ChatTongyi
from langchain_core.prompts import ChatPromptTemplate

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 示例数据：模拟 30 天、3 个渠道的 GMV 和订单数
# 后续工具都会基于这个 df 做分析
df = pd.DataFrame({
    'date': pd.date_range('2024-01-01', periods=30),
    'channel': ['A', 'B', 'C'] * 10,
    'gmv': [1000, 1200, 800, 1500, 900, 1100] * 5,
    'order_count': [50, 60, 40, 75, 45, 55] * 5
})

@tool
def inspect_schema(df_name: str) -> str:
    """检查 DataFrame 的列名和数据类型
    
    Args:
        df_name: 数据框名称（当前只支持示例数据）
    """
    # 返回 schema 信息（列名 / 维度 / dtype），帮助模型做后续决策
    return f"columns: {list(df.columns)}, shape: {df.shape}, dtypes: {df.dtypes.to_dict()}"

@tool
def compute_metrics(input_json: str) -> str:
    """按指定字段分组计算指标
    
    Args:
        input_json: JSON字符串，格式为 {"by": "字段名", "metric": "字段名", "top_n": 数字}
    """
    import json
    
    try:
        params = json.loads(input_json) if isinstance(input_json, str) else input_json
        by = params.get('by')
        metric = params.get('metric')
        top_n = params.get('top_n', 5)
    except (json.JSONDecodeError, AttributeError) as e:
        return f"参数解析错误: {e}"
    
    # 基础校验：避免模型给空字段
    if not by or not metric:
        return "错误：必须提供 by 和 metric 参数"
    
    # 字段校验：避免越权访问不存在字段
    if by not in df.columns or metric not in df.columns:
        return f"错误：字段 {by} 或 {metric} 不存在"
    
    result = df.groupby(by)[metric].sum().sort_values(ascending=False)
    return f"Top {top_n} {metric} by {by}:\n{result.head(top_n).to_string()}"

@tool
def compute_conversion_rate(input_json: str) -> str:
    """计算转化率
    
    Args:
        input_json: JSON字符串，格式为 {"by": "分组字段名"}
    """
    import json
    
    try:
        params = json.loads(input_json) if isinstance(input_json, str) else input_json
        by = params.get('by')
    except (json.JSONDecodeError, AttributeError) as e:
        return f"参数解析错误: {e}"
    
    if not by:
        return "错误：必须提供 by 参数"
    
    if by not in df.columns:
        return f"错误：字段 {by} 不存在"
    
    # 计算转化率 = 订单数 / GMV (简化示例)
    result = df.groupby(by).apply(lambda x: (x['order_count'].sum() / x['gmv'].sum()) * 100).round(2)
    return f"各{by}转化率(%):\n{result.to_string()}"

@tool
def plot_trend(input_json: str) -> str:
    """绘制趋势图并保存
    
    Args:
        input_json: JSON字符串，格式为 {"metric": "字段名", "by": "分组字段名(可选)"}
    """
    import json
    
    try:
        params = json.loads(input_json) if isinstance(input_json, str) else input_json
        metric = params.get('metric')
        by = params.get('by')
    except (json.JSONDecodeError, AttributeError) as e:
        return f"参数解析错误: {e}"
    
    if not metric:
        return "错误：必须提供 metric 参数"
    
    if metric not in df.columns:
        return f"错误：字段 {metric} 不存在"
    
    try:
        plt.figure(figsize=(10, 6))
        if by:
            if by not in df.columns:
                return f"错误：分组字段 {by} 不存在"
            for group in df[by].unique():
                data = df[df[by] == group]
                plt.plot(data['date'], data[metric], label=f'{by}={group}')
        else:
            plt.plot(df['date'], df[metric])
        
        plt.title(f'{metric} 趋势图')
        plt.xlabel('日期')
        plt.ylabel(metric)
        plt.legend()
        plt.grid(True)
        
        # 保存图片，忽略字体警告
        import warnings
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", UserWarning)  # 忽略字体警告
            plt.savefig('trend_plot.png', bbox_inches='tight')
        plt.close()
        return "趋势图已保存为 trend_plot.png"
    except Exception as e:
        return f"绘图错误: {e}"

# 初始化 LLM 和 Agent
llm = ChatTongyi(model="qwen-turbo", temperature=0)

# ReAct 提示词模板
prompt = ChatPromptTemplate.from_template("""你是一个数据分析助手。使用以下工具来回答用户的问题。

可用工具：
{tools}

工具名称：
{tool_names}

使用以下格式：
Question: 用户的问题
Thought: 我需要思考如何回答这个问题
Action: 选择一个工具
Action Input: 工具的输入参数（必须是JSON格式）
Observation: 工具的输出结果
... (可以重复 Thought/Action/Action Input/Observation 多次)
Thought: 我现在知道最终答案了
Final Answer: 基于观察结果的最终答案（不要在Final Answer后再执行任何Action）

重要提示：
- inspect_schema 工具：查看数据结构，参数：{{"df_name": "数据框名称"}}
- compute_metrics 工具：计算分组指标，参数：{{"by": "字段名", "metric": "字段名", "top_n": 数字}}
- compute_conversion_rate 工具：计算转化率，参数：{{"by": "分组字段名"}}
- plot_trend 工具：绘制趋势图，参数：{{"metric": "字段名", "by": "分组字段名(可选)"}}
- 一旦给出Final Answer，就不能再执行任何Action
- 确保在Final Answer中完成所有需要的分析

开始！

Question: {input}
Thought: {agent_scratchpad}""")

# 创建 Agent
tools = [inspect_schema, compute_metrics, compute_conversion_rate, plot_trend]
agent = create_react_agent(llm, tools, prompt)
executor = AgentExecutor(
    agent=agent, 
    tools=tools, 
    verbose=True, 
    max_iterations=5,
    handle_parsing_errors=True,  # 处理解析错误，让 Agent 重试
    return_intermediate_steps=True  # 返回中间轨迹（intermediate_steps）用于打印
)

# 执行示例
if __name__ == "__main__":
    # 测试多个工具调用
    questions = [
        "哪个渠道的GMV最高？请画出各渠道GMV的趋势图",
        "请分析各渠道的订单数量情况，并计算订单转化率",
        "查看数据结构，然后给出整体数据分析报告"
    ]
    
    for i, question in enumerate(questions, 1):
        print(f"\n{'='*60}")
        print(f"问题 {i}: {question}")
        print('='*60)
        
        result = executor.invoke({"input": question})
        output = result.get("output", None) if isinstance(result, dict) else None
        print(f"\n最终回答: {output}")
        
        # 显示工具调用统计
        steps = result.get("intermediate_steps", None) if isinstance(result, dict) else None
        if steps:
            print(f"\n工具调用轨迹:")
            for step_num, (action, observation) in enumerate(steps, 1):
                tool_name = getattr(action, "tool", None)
                tool_input = getattr(action, "tool_input", None)
                print(f"  步骤 {step_num}: 调用工具 '{tool_name}'")
                print(f"    输入参数: {tool_input}")
                obs_str = str(observation)
                print(f"    输出结果: {obs_str[:100]}..." if len(obs_str) > 100 else f"    输出结果: {obs_str}")
                print()
        else:
            print("\n工具调用轨迹: 未返回 intermediate_steps（请确认 return_intermediate_steps=True，或通过 verbose 日志查看轨迹）")
        
        print(f"{'='*60}\n")

// ============================
// 目的：给 LLM 一个“工具调用协议”的例子（结构化动作），让它输出可被系统解析的参数
//
// 含义：
// - tool：要调用的工具名（注册中心里必须存在）
// - args：该工具的入参（必须是白名单字段 + 合法范围）
//
// 注意：这里展示的是“结构化指令”，不是让模型写 Pandas 脚本
// ============================
{"tool":"groupby_sum","args":{"by":"channel","metric":"gmv","top_n":3}}

# ============================
# 目的：生产级“数据分析 Agent”骨架（可上线思路示例）
#
# 这一段强调“产品化”必备能力：
# - 用户级限流（防滥用）
# - 高危操作检测 + Human-in-the-Loop
# - 审计日志（trace_id、risk_level、耗时、状态）
# - 工具白名单 + 参数校验 + top_n/导出行数上限
# - 可观测性：日志、耗时、成功率
#
# 注意：这是教学示例，重点是架构与治理点，部分细节需按你实际框架版本做适配。
# ============================

import os
import json
import time
import logging
import uuid
import hashlib
from typing import Dict, Any, List, Optional
from datetime import datetime
from dataclasses import dataclass
from enum import Enum

from langchain.tools import tool
from langchain.agents import create_agent
from langchain.agents.middleware.model_call_limit import ModelCallLimitMiddleware
from langchain.agents.middleware.tool_call_limit import ToolCallLimitMiddleware
from langchain.agents.middleware.model_retry import ModelRetryMiddleware
from langchain.agents.middleware.tool_retry import ToolRetryMiddleware
from langchain_community.chat_models import ChatTongyi
from langchain_core.runnables import RunnableConfig
from langchain_core.messages import HumanMessage, AIMessage

# === 数据处理库 ===
import pandas as pd
import numpy as np

# === 生产环境推荐组合 ===
# 1️⃣ 安全：system_prompt + middleware（限流）+ interrupt_before（高危审批）
# 2️⃣ 可靠：middleware（重试）+ name（监控标识）  
# 3️⃣ 可观测：name + debug=False + 外部日志系统 + 审计追踪

# 配置日志系统（生产建议：输出到集中式日志平台；这里用文件+控制台演示）
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('/tmp/data_analysis_agent.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("data_analysis_agent")

# 审计日志记录器
audit_logger = logging.getLogger("audit")

# === 数据模型 ===
@dataclass
class AuditLog:
    """审计日志数据模型"""
    trace_id: str
    user_id: str
    operation: str
    status: str
    timestamp: datetime
    tools_used: List[str]
    risk_level: str
    execution_time: float
    error_message: Optional[str] = None

class RiskLevel(Enum):
    """风险级别枚举"""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"

class OperationStatus(Enum):
    """操作状态枚举"""
    PENDING = "pending"
    APPROVED = "approved"
    REJECTED = "rejected"
    CANCELLED = "cancelled"
    COMPLETED = "completed"
    FAILED = "failed"

# === 高级限流器（只保留用户维度） ===
# 说明：限流的第一层防线，按 user_id 统计调用频次，保护系统不被刷爆
class AdvancedRateLimiter:
    """高级限流器：当前示例中只做“用户级别”限流，便于和业务账号体系对齐"""
    def __init__(self):
        # 仅按 user_id 做滑动时间窗口统计
        self.user_limits: dict[str, list[float]] = {}
    
    def is_allowed(self, user_id: str) -> tuple[bool, str]:
        """检查是否允许请求，返回 (是否允许, 限制原因)

        这里的策略是：同一 user_id 在 60 秒窗口内最多 10 次调用。
        """
        now = time.time()
        key = f"user_{user_id}"
        window_seconds = 60
        max_requests = 10

        if key not in self.user_limits:
            self.user_limits[key] = []
        
        # 清理窗口外的旧请求
        self.user_limits[key] = [
            ts for ts in self.user_limits[key]
            if now - ts < window_seconds
        ]

        if len(self.user_limits[key]) >= max_requests:
            return False, f"用户限流：{user_id} 在 {window_seconds} 秒内最多 {max_requests} 次调用"

        # 记录本次调用时间戳
        self.user_limits[key].append(now)
        return True, "允许"

rate_limiter = AdvancedRateLimiter()

# === 高级安全检测器 ===
class SecurityAnalyzer:
    """高级安全分析器"""
    def __init__(self):
        self.high_risk_keywords = [
            "删除", "drop", "truncate", "清空", "覆盖", "全量", "所有数据",
            "export", "导出", "下载", "敏感", "权限", "管理员", "root"
        ]
        self.sensitive_fields = {"salary", "phone", "email", "id_card", "bank_account"}
        self.destructive_tools = {"delete_data", "truncate_table", "drop_table"}
    
    def analyze_risk(self, input_text: str, tools_used: List[str], fields_accessed: List[str]) -> RiskLevel:
        """分析操作风险级别"""
        risk_score = 0
        
        # 关键词检测
        if any(keyword in input_text.lower() for keyword in self.high_risk_keywords):
            risk_score += 3
        
        # 敏感字段检测
        if any(field in self.sensitive_fields for field in fields_accessed):
            risk_score += 2
        
        # 破坏性工具检测
        if any(tool in self.destructive_tools for tool in tools_used):
            risk_score += 3
        
        # 大量数据导出检测
        if "export" in input_text.lower() or "导出" in input_text:
            risk_score += 2
        
        # 数据量检测
        if any(word in input_text.lower() for word in ["10000", "一万", "all", "全部"]):
            risk_score += 1
        
        if risk_score >= 5:
            return RiskLevel.HIGH
        elif risk_score >= 3:
            return RiskLevel.MEDIUM
        else:
            return RiskLevel.LOW
    
    def is_high_risk_operation(self, input_text: str) -> bool:
        """检测是否为高危操作"""
        risk = self.analyze_risk(input_text, [], [])
        return risk == RiskLevel.HIGH

security_analyzer = SecurityAnalyzer()

# === 审计追踪系统 ===
class AuditTrail:
    """审计追踪系统"""
    def __init__(self):
        self.logs = []
    
    def create_log(self, user_id: str, operation: str, risk_level: RiskLevel) -> AuditLog:
        """创建审计日志"""
        trace_id = str(uuid.uuid4())[:8]
        return AuditLog(
            trace_id=trace_id,
            user_id=user_id,
            operation=operation,
            status=OperationStatus.PENDING.value,
            timestamp=datetime.now(),
            tools_used=[],
            risk_level=risk_level.value,
            execution_time=0.0
        )
    
    def log_operation(self, audit_log: AuditLog):
        """记录操作日志"""
        log_data = {
            "trace_id": audit_log.trace_id,
            "user_id": audit_log.user_id,
            "operation": audit_log.operation,
            "status": audit_log.status,
            "timestamp": audit_log.timestamp.isoformat(),
            "tools_used": audit_log.tools_used,
            "risk_level": audit_log.risk_level,
            "execution_time": audit_log.execution_time,
            "error_message": audit_log.error_message
        }
        
        # 记录到审计日志
        audit_logger.info(json.dumps(log_data, ensure_ascii=False))
        
        # 记录到内存（实际生产环境应该用数据库）
        self.logs.append(log_data)

audit_trail = AuditTrail()

# === 性能监控器 ===
class PerformanceMonitor:
    """性能监控器"""
    def __init__(self):
        self.metrics = {}
    
    def record_execution(self, tool_name: str, execution_time: float, success: bool):
        """记录执行指标"""
        if tool_name not in self.metrics:
            self.metrics[tool_name] = {
                "total_calls": 0,
                "success_calls": 0,
                "total_time": 0.0,
                "avg_time": 0.0,
                "success_rate": 0.0
            }
        
        metrics = self.metrics[tool_name]
        metrics["total_calls"] += 1
        metrics["total_time"] += execution_time
        metrics["avg_time"] = metrics["total_time"] / metrics["total_calls"]
        
        if success:
            metrics["success_calls"] += 1
        
        metrics["success_rate"] = metrics["success_calls"] / metrics["total_calls"]
        
        # 记录到日志
        logger.info(f"Performance: {tool_name} - avg_time: {metrics['avg_time']:.3f}s, success_rate: {metrics['success_rate']:.2%}")

performance_monitor = PerformanceMonitor()


# === 示例数据 ===
import pandas as pd
import numpy as np

# 生成更丰富的示例数据
np.random.seed(42)
dates = pd.date_range('2024-01-01', periods=100)
channels = ['A', 'B', 'C', 'D', 'E']
regions = ['北京', '上海', '广州', '深圳', '杭州']

df = pd.DataFrame({
    'date': np.random.choice(dates, 1000),
    'channel': np.random.choice(channels, 1000),
    'region': np.random.choice(regions, 1000),
    'gmv': np.random.normal(1000, 300, 1000).astype(int),
    'order_count': np.random.poisson(50, 1000),
    'user_count': np.random.poisson(100, 1000),
    'product_category': np.random.choice(['电子产品', '服装', '食品', '家居'], 1000)
})

# 确保数据合理性
df['gmv'] = df['gmv'].clip(lower=100)
df['order_count'] = df['order_count'].clip(lower=1)
df['user_count'] = df['user_count'].clip(lower=1)

# === 数据分析工具（生产级封装） ===
# 说明：工具化的核心是“白名单 + 上限”，确保模型只能在安全边界内做分析
ALLOWED_COLUMNS = {"channel", "gmv", "order_count", "date", "region", "user_count", "product_category"}
ALLOWED_AGGREGATIONS = {"sum", "mean", "count", "max", "min", "std"}
MAX_TOP_N = 100
MAX_EXPORT_ROWS = 10000

@tool
def safe_groupby(input_json: str) -> str:
    """安全的分组聚合（生产级实现）
    
    Args:
        input_json: JSON格式，{"by": ["字段名"], "metric": "字段名", "agg": "聚合方式", "top_n": 数字}
    """
    # 记录耗时，用于性能追踪
    start_time = time.time()
    tool_name = "safe_groupby"
    
    try:
        # 解析参数（来自 LLM 的结构化 JSON）
        params = json.loads(input_json)
        group_by = params.get("by", [])
        metric = params.get("metric")
        agg_func = params.get("agg", "sum")
        top_n = min(params.get("top_n", 5), MAX_TOP_N)
        
        # 参数验证（字段白名单 + 聚合方式白名单 + top_n 上限）
        if not isinstance(group_by, list):
            group_by = [group_by]
        
        if not all(col in ALLOWED_COLUMNS for col in group_by + [metric]):
            return f"错误：只允许字段 {ALLOWED_COLUMNS}"
        
        if agg_func not in ALLOWED_AGGREGATIONS:
            return f"错误：只允许聚合方式 {ALLOWED_AGGREGATIONS}"
        
        # 执行聚合（只允许预定义操作，避免复杂/危险表达式）
        if agg_func == "sum":
            result = df.groupby(group_by)[metric].sum().nlargest(top_n)
        elif agg_func == "mean":
            result = df.groupby(group_by)[metric].mean().nlargest(top_n)
        elif agg_func == "count":
            result = df.groupby(group_by)[metric].count().nlargest(top_n)
        elif agg_func == "max":
            result = df.groupby(group_by)[metric].max().nlargest(top_n)
        elif agg_func == "min":
            result = df.groupby(group_by)[metric].min().nlargest(top_n)
        elif agg_func == "std":
            result = df.groupby(group_by)[metric].std().nlargest(top_n)
        
        result_str = result.to_string()
        
        # 记录性能指标
        execution_time = time.time() - start_time
        performance_monitor.record_execution(tool_name, execution_time, True)
        
        # 监控记录
        logger.info(json.dumps({
            "timestamp": datetime.now().isoformat(),
            "tool_name": tool_name,
            "execution_time": round(execution_time, 3),
            "status": "success",
            "params": input_json[:100] + "..." if len(input_json) > 100 else input_json,
            "result_rows": len(result)
        }))
        
        return result_str
        
    except Exception as e:
        execution_time = time.time() - start_time
        performance_monitor.record_execution(tool_name, execution_time, False)
        
        logger.error(json.dumps({
            "timestamp": datetime.now().isoformat(),
            "tool_name": tool_name,
            "execution_time": round(execution_time, 3),
            "status": "error",
            "error": str(e),
            "params": input_json[:100] + "..." if len(input_json) > 100 else input_json
        }))
        
        return f"执行错误：{str(e)}"

@tool
def safe_filter(input_json: str) -> str:
    """安全的数据过滤
    
    Args:
        input_json: JSON格式，{"conditions": [{"field": "字段名", "operator": "操作符", "value": "值"}], "limit": 数字}
    """
    start_time = time.time()
    tool_name = "safe_filter"
    
    try:
        params = json.loads(input_json)
        conditions = params.get("conditions", [])
        limit = min(params.get("limit", 100), MAX_EXPORT_ROWS)
        
        # 参数验证
        for condition in conditions:
            field = condition.get("field")
            if field not in ALLOWED_COLUMNS:
                return f"错误：不允许字段 {field}"
        
        # 构建过滤条件
        filtered_df = df.copy()
        for condition in conditions:
            field = condition["field"]
            operator = condition["operator"]
            value = condition["value"]
            
            if operator == "eq":
                filtered_df = filtered_df[filtered_df[field] == value]
            elif operator == "gt":
                filtered_df = filtered_df[filtered_df[field] > value]
            elif operator == "lt":
                filtered_df = filtered_df[filtered_df[field] < value]
            elif operator == "in":
                if isinstance(value, list):
                    filtered_df = filtered_df[filtered_df[field].isin(value)]
                else:
                    filtered_df = filtered_df[filtered_df[field] == value]
        
        result = filtered_df.head(limit).to_string()
        
        # 记录性能指标
        execution_time = time.time() - start_time
        performance_monitor.record_execution(tool_name, execution_time, True)
        
        return result
        
    except Exception as e:
        execution_time = time.time() - start_time
        performance_monitor.record_execution(tool_name, execution_time, False)
        return f"过滤错误：{str(e)}"

@tool
def detect_anomalies(input_json: str) -> str:
    """异常检测
    
    Args:
        input_json: JSON格式，{"metric": "指标名", "method": "检测方法", "threshold": 阈值}
    """
    start_time = time.time()
    tool_name = "detect_anomalies"
    
    try:
        params = json.loads(input_json)
        metric = params.get("metric")
        method = params.get("method", "zscore")
        threshold = params.get("threshold", 2.0)
        
        if metric not in ALLOWED_COLUMNS:
            return f"错误：不允许字段 {metric}"
        
        # 简单的异常检测
        if method == "zscore":
            z_scores = np.abs((df[metric] - df[metric].mean()) / df[metric].std())
            anomalies = df[z_scores > threshold]
        
        result = f"检测到 {len(anomalies)} 个异常点\n"
        result += f"异常值范围：{anomalies[metric].min():.2f} - {anomalies[metric].max():.2f}\n"
        result += "异常样本：\n" + anomalies.head(5).to_string()
        
        # 记录性能指标
        execution_time = time.time() - start_time
        performance_monitor.record_execution(tool_name, execution_time, True)
        
        return result
        
    except Exception as e:
        execution_time = time.time() - start_time
        performance_monitor.record_execution(tool_name, execution_time, False)
        return f"异常检测错误：{str(e)}"

# === 创建生产级 Agent ===
# 说明：create_agent 负责把 model + tools + system_prompt 组装为一个可调用的 Agent
api_key = os.getenv("DASHSCOPE_API_KEY")
llm = ChatTongyi(model_name="qwen-turbo", dashscope_api_key=api_key, temperature=0)

# 系统提示词（安全边界）：把“不能做什么/只能做什么”写清楚，比事后兜底更重要
SYSTEM_PROMPT = """你是一个智能数据分析助手。请遵循以下规则：

🔒 安全约束：
1. 只能使用提供的工具，不能执行任意代码
2. 敏感操作（删除、导出全量数据等）需要人工审批
3. 严格遵守数据隐私和权限边界

📊 分析规范：
1. 先理解业务问题，再选择合适的分析工具
2. 确保分析结果可解释、可复现
3. 重要结论要提供数据支撑

⚡ 性能要求：
1. 优先使用预聚合数据，避免全表扫描
2. 合理设置 top_n，防止返回过大数据集
3. 复杂分析要分步骤执行

🛠️ 可用工具：
- safe_groupby: 分组聚合分析
- safe_filter: 数据过滤查询
- detect_anomalies: 异常检测分析"""

# 创建 Agent（带监控标识）
# name：用于监控标识（例如打到 tracing/metrics 系统）
data_analysis_agent = create_agent(
    model=llm,
    tools=[safe_groupby, safe_filter, detect_anomalies],
    system_prompt=SYSTEM_PROMPT,
    name="data_analysis_agent",  # 监控标识
    debug=False,  # 生产环境关闭 debug
    # 中间件：示例展示按“单次运行”维度控制调用次数
    # 注意：ModelCallLimitMiddleware / ToolCallLimitMiddleware 需根据你实际使用的 LangChain 版本
    # 引入对应实现或自定义实现，这里只作为 API 用法示意。
    middleware=(
        # 限流保护：防止单次运行内滥用模型 / 工具
        ModelCallLimitMiddleware(run_limit=10),    # 每次运行最多 10 次模型调用
        ToolCallLimitMiddleware(run_limit=5),      # 每次运行最多 5 次工具调用
    ),
)

# === 高级生产级调用包装器 ===
class ProductionAgentWrapper:
    """生产级 Agent 包装器：完整的安全、监控、审计功能"""
    
    def __init__(self, agent, user_id: str = "default_user"):
        self.agent = agent
        self.user_id = user_id
    
    def invoke(self, user_input: str, require_approval: bool = True) -> Dict[str, Any]:
        """生产级 Agent 调用"""
        trace_id = str(uuid.uuid4())[:8]
        start_time = time.time()
        
        # 创建审计日志
        risk_level = security_analyzer.analyze_risk(user_input, [], [])
        audit_log = audit_trail.create_log(self.user_id, user_input, risk_level)
        
        try:
            # 1. 用户级限流检查（第一道保护）
            allowed, reason = rate_limiter.is_allowed(self.user_id)
            if not allowed:
                audit_log.status = OperationStatus.REJECTED.value
                audit_log.error_message = reason
                audit_trail.log_operation(audit_log)
                
                return {
                    "status": "rate_limited",
                    "message": reason,
                    "trace_id": trace_id,
                    "agent_name": "data_analysis_agent"
                }
            
            # 2. 高危操作检测（需要人工确认时先返回 pending）
            if require_approval and risk_level == RiskLevel.HIGH:
                audit_log.status = OperationStatus.PENDING.value
                audit_trail.log_operation(audit_log)
                
                return {
                    "status": "pending_approval",
                    "message": "检测到高危操作，需要人工审批",
                    "requires_approval": True,
                    "operation": user_input,
                    "trace_id": trace_id,
                    "risk_level": risk_level.value,
                    "agent_name": "data_analysis_agent"
                }
            
            # 3. 执行分析：将用户输入封装为 messages 交给 Agent
            config = RunnableConfig(
                callbacks=[],
                tags=["production", "data_analysis"],
                metadata={"trace_id": trace_id, "user_id": self.user_id}
            )
            
            result = self.agent.invoke(
                {"messages": [HumanMessage(content=user_input)]},
                config=config
            )
            
            # 4. 记录成功日志
            execution_time = time.time() - start_time
            audit_log.status = OperationStatus.COMPLETED.value
            audit_log.execution_time = execution_time
            audit_log.tools_used = ["data_analysis_agent"]
            audit_trail.log_operation(audit_log)
            
            # 5. 记录性能指标
            logger.info(json.dumps({
                "timestamp": datetime.now().isoformat(),
                "trace_id": trace_id,
                "user_id": self.user_id,
                "agent_name": "data_analysis_agent",
                "execution_time": round(execution_time, 3),
                "status": "success",
                "query": user_input[:100] + "..." if len(user_input) > 100 else user_input,
                "risk_level": risk_level.value
            }))
            
            return {
                "status": "success",
                "result": result["messages"][-1].content,
                "trace_id": trace_id,
                "agent_name": "data_analysis_agent",
                "execution_time": round(execution_time, 3),
                "risk_level": risk_level.value
            }
            
        except Exception as e:
            execution_time = time.time() - start_time
            
            # 记录错误日志
            audit_log.status = OperationStatus.FAILED.value
            audit_log.execution_time = execution_time
            audit_log.error_message = str(e)
            audit_trail.log_operation(audit_log)
            
            logger.error(json.dumps({
                "timestamp": datetime.now().isoformat(),
                "trace_id": trace_id,
                "user_id": self.user_id,
                "agent_name": "data_analysis_agent",
                "execution_time": round(execution_time, 3),
                "status": "error",
                "error": str(e),
                "query": user_input
            }))
            
            return {
                "status": "error",
                "error": str(e),
                "trace_id": trace_id,
                "agent_name": "data_analysis_agent"
            }
    
    def get_performance_metrics(self) -> Dict[str, Any]:
        """获取性能指标"""
        return {
            "performance_metrics": performance_monitor.metrics,
            "recent_logs": audit_trail.logs[-10:],  # 最近10条日志
            "system_status": "healthy"
        }

# 创建生产级 Agent 实例
production_agent = ProductionAgentWrapper(data_analysis_agent, "user_001")

# === 使用示例 ===
if __name__ == "__main__":
    print("=== 生产级数据分析系统演示 ===\n")
    
    # 示例1：正常分析请求
    print("1. 正常分析请求：")
    response1 = production_agent.invoke(
        '{"by": ["channel"], "metric": "gmv", "agg": "sum", "top_n": 3}'
    )
    print(f"状态: {response1['status']}")
    print(f"结果: {response1.get('result', '无结果')[:100]}...")
    print(f"执行时间: {response1.get('execution_time', 0)}秒")
    print(f"风险级别: {response1.get('risk_level', 'unknown')}")
    print()
    
    # 示例2：数据过滤请求
    print("2. 数据过滤请求：")
    response2 = production_agent.invoke(
        '{"conditions": [{"field": "channel", "operator": "in", "value": ["A", "B"]}], "limit": 5}'
    )
    print(f"状态: {response2['status']}")
    print(f"结果预览: {response2.get('result', '无结果')[:100]}...")
    print()
    
    # 示例3：异常检测请求
    print("3. 异常检测请求：")
    response3 = production_agent.invoke(
        '{"metric": "gmv", "method": "zscore", "threshold": 2.0}'
    )
    print(f"状态: {response3['status']}")
    print(f"结果预览: {response3.get('result', '无结果')[:100]}...")
    print()
    
    # 示例4：高危操作请求
    print("4. 高危操作请求：")
    response4 = production_agent.invoke("删除所有渠道数据")
    print(f"状态: {response4['status']}")
    print(f"消息: {response4.get('message', '无消息')}")
    print(f"需要审批: {response4.get('requires_approval', False)}")
    print()
    
    # 示例5：获取性能指标
    print("5. 性能指标：")
    metrics = production_agent.get_performance_metrics()
    print(f"工具调用统计: {len(metrics['performance_metrics'])} 个工具")
    for tool_name, tool_metrics in metrics['performance_metrics'].items():
        print(f"  {tool_name}: {tool_metrics['total_calls']}次调用, "
              f"成功率 {tool_metrics['success_rate']:.1%}, "
              f"平均耗时 {tool_metrics['avg_time']:.3f}秒")
    print()
    
    print("=== 演示完成 ===")

# ============================
# 目的：展示“先计划再执行”的最小闭环（Plan-and-Execute 思路）
#
# 你将看到：
# - 第一步让模型把自然语言问题“翻译”为结构化计划（JSON）
# - 第二步由系统按计划逐步调用工具（或函数）执行
# - 计划受 schema 约束，保证口径一致、可复现、可审计
#
# 注意：这段代码用通义千问（ChatTongyi）演示“真实 API 调用”，并保持与本章其它模块一致。
# 安全建议：
# - API Key 必须从环境变量读取，禁止写死在代码里
# - 计划必须做白名单约束（允许工具、允许字段、最大步数）
# ============================

import json
import time
import uuid
import os
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple
from langchain_community.chat_models import ChatTongyi
from langchain_core.messages import SystemMessage, HumanMessage

# 初始化大模型（与本章其它代码保持一致）
api_key = os.getenv("DASHSCOPE_API_KEY")
if not api_key:
    raise ValueError("请先在环境变量中设置 DASHSCOPE_API_KEY")
llm = ChatTongyi(model_name="qwen-turbo", dashscope_api_key=api_key, temperature=0)

# 模拟业务数据
sample_data = {
    "sales_data": [
        {"date": "2024-01-01", "channel": "A", "gmv": 1200, "orders": 60},
        {"date": "2024-01-02", "channel": "B", "gmv": 800, "orders": 40},
        {"date": "2024-01-03", "channel": "C", "gmv": 1500, "orders": 75},
        # ... 更多数据
    ]
}

# 计划 Schema：约束模型输出，避免“随口编计划”
PLAN_SCHEMA = {
    "required": ["intent", "range", "steps", "deliver"],
    "step_required": ["tool", "args"],
    "allowed_tools": {"safe_filter", "safe_groupby", "detect_anomalies", "plot"},
    "allowed_ranges": {"last_7d", "last_30d", "this_month", "custom"},
    "max_steps": 6,
}

# 增强的计划生成提示词：让模型理解业务上下文，并输出“可执行计划”
PLAN_PROMPT = """你是专业的数据分析规划师。请将用户问题转换为结构化的分析计划JSON。

业务背景：
- 数据源：电商销售数据，包含字段：date(日期), channel(渠道), gmv(销售额), orders(订单数)
- 支持的分析：趋势分析、渠道对比、异常检测、转化率计算
- 输出要求：可执行的工具链步骤

请严格按照以下JSON格式输出，不要包含任何其他文字：
{
  "intent": "分析意图描述",
  "range": "时间范围(last_7d/last_30d/this_month/custom)",
  "steps": [
    {"tool": "工具名", "args": {"参数名": "参数值"}},
    ...
  ],
  "deliver": "交付物描述"
}

可用工具：
- safe_filter: 数据过滤 {"conditions": [{"field": "字段", "op": ">", "value": "数值"}], "limit": 数量}
- safe_groupby: 分组统计 {"by": ["字段"], "metric": "指标", "agg": "sum/avg/count", "top_n": 数量}
- detect_anomalies: 异常检测 {"metric": "指标", "method": "std/iqr"}
- plot: 图表生成 {"chart": "line/bar/pie", "x": "X轴字段", "y": "Y轴字段", "series": "分组字段"}

示例：
用户问题："近30天各渠道GMV趋势分析"
输出：
{
  "intent": "分析近30天各渠道GMV趋势变化",
  "range": "last_30d",
  "steps": [
    {"tool": "safe_filter", "args": {"conditions": [{"field": "date", "op": ">=", "value": "2024-01-01"}], "limit": 1000}},
    {"tool": "safe_groupby", "args": {"by": ["channel", "date"], "metric": "gmv", "agg": "sum", "top_n": 50}},
    {"tool": "plot", "args": {"chart": "line", "x": "date", "y": "gmv", "series": "channel"}}
  ],
  "deliver": "渠道GMV趋势图+数据洞察"
}

现在请处理用户问题："""

@dataclass
class TraceEvent:
    name: str
    ok: bool
    ms: int
    meta: Dict[str, Any]

class Trace:
    def __init__(self, query: str):
        self.trace_id = f"t-{uuid.uuid4().hex[:6]}"
        self.query = query
        self.start = time.time()
        self.events: List[TraceEvent] = []

    def log(self, name: str, ok: bool, meta: Optional[Dict[str, Any]] = None):
        meta = meta or {}
        self.events.append(
            TraceEvent(
                name=name,
                ok=ok,
                ms=round((time.time() - self.start) * 1000),
                meta=meta,
            )
        )

    def summary(self) -> Dict[str, Any]:
        return {
            "trace_id": self.trace_id,
            "query": self.query,
            "events": [e.__dict__ for e in self.events],
        }

def call_llm_for_planning(query: str) -> Dict[str, Any]:
    """使用真实大模型生成分析计划"""
    try:
        resp = llm.invoke(
            [
                SystemMessage(content="你是专业的数据分析规划师，只输出JSON格式的分析计划。"),
                HumanMessage(content=PLAN_PROMPT + query),
            ]
        )
        content = (resp.content or "").strip()
        
        # 清理可能的markdown标记
        if content.startswith("```json"):
            content = content[7:]
        if content.endswith("```"):
            content = content[:-3]
        
        return json.loads(content)
    except Exception as e:
        print(f"大模型调用失败: {e}")
        # 降级到示例计划
        return {
            "intent": "分析渠道GMV趋势",
            "range": "last_30d",
            "steps": [
                {"tool": "safe_filter", "args": {"conditions": [], "limit": 1000}},
                {"tool": "safe_groupby", "args": {"by": ["channel"], "metric": "gmv", "agg": "sum", "top_n": 5}},
                {"tool": "plot", "args": {"chart": "line", "x": "date", "y": "gmv", "series": "channel"}}
            ],
            "deliver": "图表+结论"
        }

def parse_query_to_plan(query: str, use_llm: bool = True) -> Dict[str, Any]:
    """自然语言 → 结构化分析计划"""
    if use_llm:
        return call_llm_for_planning(query)
    else:
        # 直接执行模式（无计划）
        return {"direct_execution": True, "query": query}

def validate_plan(plan: Dict[str, Any]) -> Tuple[bool, str]:
    """计划校验：必填字段、工具白名单、步数上限、字段类型"""
    if "direct_execution" in plan:
        return True, "直接执行模式"
    
    missing = [k for k in PLAN_SCHEMA["required"] if k not in plan]
    if missing:
        return False, f"计划缺少字段: {missing}"

    if plan.get("range") not in PLAN_SCHEMA["allowed_ranges"]:
        return False, f"range 不合法: {plan.get('range')}"

    steps = plan.get("steps")
    if not isinstance(steps, list) or len(steps) == 0:
        return False, "steps 必须为非空数组"
    if len(steps) > PLAN_SCHEMA["max_steps"]:
        return False, f"steps 过多: {len(steps)} > {PLAN_SCHEMA['max_steps']}"

    for i, s in enumerate(steps):
        if not isinstance(s, dict):
            return False, f"step[{i}] 必须为对象"
        for k in PLAN_SCHEMA["step_required"]:
            if k not in s:
                return False, f"step[{i}] 缺少字段: {k}"
        if s["tool"] not in PLAN_SCHEMA["allowed_tools"]:
            return False, f"step[{i}] tool 不允许: {s['tool']}"
        if not isinstance(s["args"], dict):
            return False, f"step[{i}] args 必须为对象"

    return True, "ok"

def assess_risk(plan: Dict[str, Any]) -> str:
    """风险评估"""
    if "direct_execution" in plan:
        return "high"  # 直接执行风险较高
    
    HIGH_RISK = {"export", "delete", "drop", "truncate", "salary", "id_card"}
    text = json.dumps(plan, ensure_ascii=False).lower()
    return "high" if any(k in text for k in HIGH_RISK) else "low"

def compile_plan(plan: Dict[str, Any]) -> List[Tuple[str, Dict[str, Any]]]:
    """计划 → 可执行步骤"""
    if "direct_execution" in plan:
        return [("direct_execute", {"query": plan["query"]})]
    return [(s["tool"], s["args"]) for s in plan["steps"]]

def execute_plan(steps: List[Tuple[str, Dict[str, Any]]], tools: Dict[str, Any], trace: Trace) -> Dict[str, Any]:
    """按步骤执行：每步记录 trace"""
    ctx: Dict[str, Any] = {}
    for tool_name, args in steps:
        t0 = time.time()
        try:
            fn = tools[tool_name]
            out = fn(args, ctx)
            ctx[tool_name] = out
            trace.log(tool_name, True, {"ms_step": round((time.time() - t0) * 1000), "result_count": len(str(out))})
        except Exception as e:
            trace.log(tool_name, False, {"error": str(e)})
            raise
    return ctx

# 工具实现
def safe_filter_tool(args: Dict[str, Any], ctx: Dict[str, Any]) -> Dict[str, Any]:
    """安全的数据过滤工具"""
    print(f"🔍 执行数据过滤: {args}")
    time.sleep(0.1)  # 模拟处理时间
    return {"filtered_rows": 856, "data_preview": "date,channel,gmv,orders"}

def safe_groupby_tool(args: Dict[str, Any], ctx: Dict[str, Any]) -> Dict[str, Any]:
    """安全的分组统计工具"""
    print(f"📊 执行分组统计: {args}")
    time.sleep(0.2)  # 模拟计算时间
    return {
        "groupby_result": [
            {"channel": "A", "gmv": 15600, "orders": 780},
            {"channel": "B", "gmv": 12300, "orders": 615},
            {"channel": "C", "gmv": 18900, "orders": 945}
        ]
    }

def detect_anomalies_tool(args: Dict[str, Any], ctx: Dict[str, Any]) -> Dict[str, Any]:
    """异常检测工具"""
    print(f"🚨 执行异常检测: {args}")
    time.sleep(0.15)  # 模拟检测时间
    return {"anomalies": 2, "anomaly_points": ["2024-01-15", "2024-01-23"]}

def plot_tool(args: Dict[str, Any], ctx: Dict[str, Any]) -> str:
    """图表生成工具"""
    print(f"📈 生成图表: {args}")
    time.sleep(0.3)  # 模拟绘图时间
    return "chart_saved_to=/output/channel_gmv_trend.png"

def call_llm_direct_answer(query: str) -> str:
    """直接执行模式：一次大模型调用直出答案（更贴近真实世界的“无计划”形态）"""
    resp = llm.invoke(
        [
            SystemMessage(content=(
                "你是数据分析助手。现在是‘直接回答’模式："
                "不要输出JSON计划，不要输出工具步骤，只输出最终结论与关键依据。"
                "如果缺少数据，请明确说明假设与不确定性。"
            )),
            HumanMessage(content=query),
        ]
    )
    return (resp.content or "").strip()

def direct_execute_tool(args: Dict[str, Any], ctx: Dict[str, Any]) -> Dict[str, Any]:
    """直接执行工具（无计划模式）"""
    query = args["query"]
    print(f"⚡ 直接执行（一次 LLM 调用直出）: {query}")
    answer = call_llm_direct_answer(query)
    return {"answer": answer}

def demonstrate_planning_vs_direct():
    """演示：先计划再执行 vs 直接执行的效果对比"""
    
    test_queries = [
        "近30天各渠道GMV趋势分析，找出异常点",
        "对比不同渠道的订单转化率，给出优化建议",
        "分析上周销售数据，识别表现异常的产品"
    ]
    
    tools = {
        "safe_filter": safe_filter_tool,
        "safe_groupby": safe_groupby_tool,
        "detect_anomalies": detect_anomalies_tool,
        "plot": plot_tool,
        "direct_execute": direct_execute_tool,
    }
    
    print("=" * 80)
    print("🎯 先结构化计划，再执行 - vs - 直接执行模式对比")
    print("=" * 80)
    
    for i, query in enumerate(test_queries, 1):
        print(f"\n📋 测试问题 {i}: {query}")
        print("-" * 60)
        
        # 模式1：先计划再执行
        print("\n🔧 模式1: 先结构化计划，再执行")
        trace1 = Trace(query)
        
        try:
            plan = parse_query_to_plan(query, use_llm=True)
            print(f"📝 生成的计划:")
            print(json.dumps(plan, ensure_ascii=False, indent=2))
            
            ok, reason = validate_plan(plan)
            trace1.log("validate_plan", ok, {"reason": reason})
            print(f"✅ 计划验证: {'通过' if ok else '失败'} - {reason}")
            
            if ok:
                risk = assess_risk(plan)
                trace1.log("assess_risk", True, {"risk": risk})
                print(f"⚠️ 风险评估: {risk}")
                
                steps = compile_plan(plan)
                print(f"🎯 执行步骤: {len(steps)} 步")
                
                ctx = execute_plan(steps, tools, trace1)
                print(f"🎉 执行完成，结果: {len(ctx)} 个工具输出")
            else:
                print("❌ 计划验证失败，跳过执行")
                
        except Exception as e:
            print(f"❌ 计划模式执行失败: {e}")
        
        print(f"\n📊 计划模式执行轨迹:")
        for event in trace1.events:
            status = "✅" if event.ok else "❌"
            print(f"  {status} {event.name} ({event.ms}ms) - {event.meta}")
        
        # 模式2：直接执行
        print("\n⚡ 模式2: 直接执行（无计划）")
        trace2 = Trace(query)
        
        try:
            plan2 = parse_query_to_plan(query, use_llm=False)
            steps2 = compile_plan(plan2)
            ctx2 = execute_plan(steps2, tools, trace2)
            print(f"🎉 直接执行完成")
            if isinstance(ctx2, dict) and "direct_execute" in ctx2 and isinstance(ctx2["direct_execute"], dict):
                ans = ctx2["direct_execute"].get("answer", "")
                if ans:
                    print("\n🧾 直接模式答案预览：")
                    print(ans[:400] + ("..." if len(ans) > 400 else ""))
            
        except Exception as e:
            print(f"❌ 直接执行失败: {e}")
        
        print(f"\n📊 直接执行轨迹:")
        for event in trace2.events:
            status = "✅" if event.ok else "❌"
            print(f"  {status} {event.name} ({event.ms}ms) - {event.meta}")
        
        # 对比总结
        print(f"\n🔄 效果对比:")
        plan_steps = len([e for e in trace1.events if e.name in ["safe_filter", "safe_groupby", "detect_anomalies", "plot"]])
        direct_steps = len([e for e in trace2.events if e.name == "direct_execute"])
        
        plan_time = sum(e.ms for e in trace1.events)
        direct_time = sum(e.ms for e in trace2.events)
        
        print(f"  📈 计划模式: {plan_steps} 个工具调用, 耗时 {plan_time}ms")
        print(f"  ⚡ 直接模式: {direct_steps} 个工具调用, 耗时 {direct_time}ms")
        print(f"  🎯 计划模式优势: 可追溯、可复现、可控制、可审计")
        print(f"  ⚡ 直接模式优势: 响应快、简单直接")
        print(f"  📝 建议: 生产环境推荐计划模式，开发调试可用直接模式")

if __name__ == "__main__":
    # 运行演示
    demonstrate_planning_vs_direct()
    
    print("\n" + "=" * 80)
    print("🎯 核心结论：先结构化计划，再执行的价值")
    print("=" * 80)
    print("✅ 可复现性: 同一问题多次执行，步骤和结果一致")
    print("✅ 可控制性: 每个步骤都有明确的输入输出，便于调试")
    print("✅ 可审计性: 完整的执行轨迹，满足合规要求")
    print("✅ 可扩展性: 新增工具只需扩展计划模板")
    print("✅ 可解释性: 用户能理解AI的\"思考过程\"")
    print("⚠️ 权衡: 增加了执行时间和复杂度")
    print("🚀 生产建议: 关键业务必须使用计划模式，简单查询可考虑直接模式")

// ============================
// 目的：审计日志（audit log）示例 —— “拒绝也要留痕”
//
// 字段解释：
// - trace_id：一次请求/一次任务的全链路追踪 ID
// - user：触发操作的用户（或角色）
// - tool：尝试调用的工具
// - args：工具入参（必要时需要脱敏/截断）
// - status：allow / deny / pending 等
// - reason：拒绝原因（权限不足、超限、风险过高...）
// ============================
{"trace_id":"t-xxx","user":"ops","tool":"export_report","args":{"format":"pdf"},"status":"deny","reason":"permission"}

# ============================
# 目的：最小“权限 / 审计 / 限流”可运行示例（复制即可跑）
#
# 你将得到：
# - 工具注册（工具名/权限级别/实现函数）
# - 权限校验：deny 也必须写审计
# - 用户维度限流：同一用户 1 秒内最多 N 次
# - 统一的受保护调用入口：guarded_call(...)
# ============================

import json
import time
import uuid
from collections import deque
from typing import Any, Callable, Dict, Deque


# 角色等级映射：把“角色字符串”转成可比较的数字。
# 生产环境通常来自：账号系统 / IAM / OAuth claims。
ROLE_LEVELS = {"viewer": 1, "editor": 2, "admin": 3}


def audit_log(
    *,
    trace_id: str,
    user: str,
    role: str,
    tool: str,
    args: Dict[str, Any],
    status: str,
    reason: str = "",
    ms: int = 0,
):
    # 审计日志：建议在生产写入数据库 / ES / 日志平台。
    # 这里为了“复制即可跑”，直接 print JSON（一行一条，便于 grep）。
    record = {
        "trace_id": trace_id,
        "user": user,
        "role": role,
        "tool": tool,
        "args": args,
        "status": status,
        "reason": reason,
        "ms": ms,
        "ts": round(time.time(), 3),
    }
    print(json.dumps(record, ensure_ascii=False))


class UserRateLimiter:
    # 用户维度限流（滑动窗口）：
    # - window_seconds：窗口大小（秒）
    # - max_calls：窗口内允许的最大调用数
    # 说明：这是“系统保护”的第一道门（防刷、防爆）。
    def __init__(self, max_calls: int, window_seconds: float):
        self.max_calls = max_calls
        self.window_seconds = window_seconds
        self.calls: Dict[str, Deque[float]] = {}

    def allow(self, user: str) -> bool:
        # deque 里存该 user 的历史调用时间戳；每次调用先淘汰窗口外的数据。
        now = time.time()
        q = self.calls.setdefault(user, deque())
        while q and now - q[0] > self.window_seconds:
            q.popleft()
        if len(q) >= self.max_calls:
            return False
        q.append(now)
        return True


def inspect_schema(_: Dict[str, Any]) -> Dict[str, Any]:
    # 示例“低风险工具”：查看 schema。
    # 要点：工具函数签名统一为 fn(args) -> json-serializable 输出。
    return {"columns": ["date", "channel", "gmv", "orders"], "rows": 1000}


def export_report(args: Dict[str, Any]) -> Dict[str, Any]:
    # 示例“高风险工具”：导出报告。
    # 在企业里 export 往往涉及：大范围数据 + 敏感字段 + 外发链路。
    fmt = args.get("format", "csv")
    return {"export": "ok", "format": fmt, "url": "https://example.com/report"}


# 工具注册中心（最小版）：
# - level：调用该工具所需的最低角色等级
# - fn：工具实现
# 生产环境可扩展：参数 schema、超时、计费、owner、风险等级、脱敏策略等。
TOOL_REGISTRY: Dict[str, Dict[str, Any]] = {
    "inspect_schema": {"level": "viewer", "fn": inspect_schema},
    "export_report": {"level": "admin", "fn": export_report},
}


def check_permission(role: str, tool_name: str) -> bool:
    # 权限校验：
    # - tool 不存在：直接拒绝
    # - role_level >= tool_required_level：允许
    tool = TOOL_REGISTRY.get(tool_name)
    if not tool:
        return False
    return ROLE_LEVELS.get(role, 0) >= ROLE_LEVELS.get(tool["level"], 999)


# 限流器：这里设置为 1 秒内最多 3 次，仅用于演示。
# 实战里你会按：用户等级/接口/工具/全局分别限流，且通常配合 Redis。
rate_limiter = UserRateLimiter(max_calls=3, window_seconds=1.0)


def guarded_call(user: str, role: str, tool: str, args: Dict[str, Any]) -> Dict[str, Any]:
    # 统一的“受保护工具调用入口”（治理核心）：
    # 1) 生成 trace_id（全链路追踪）
    # 2) 限流：超限直接 deny + 审计
    # 3) 鉴权：无权限直接 deny + 审计
    # 4) 执行工具：allow + 审计（含耗时）
    trace_id = f"t-{uuid.uuid4().hex[:6]}"
    t0 = time.time()

    if not rate_limiter.allow(user):
        audit_log(trace_id=trace_id, user=user, role=role, tool=tool, args=args, status="deny", reason="rate_limit")
        return {"ok": False, "error": "rate_limited", "trace_id": trace_id}

    if not check_permission(role, tool):
        audit_log(trace_id=trace_id, user=user, role=role, tool=tool, args=args, status="deny", reason="permission")
        return {"ok": False, "error": "permission_denied", "trace_id": trace_id}

    fn: Callable[[Dict[str, Any]], Any] = TOOL_REGISTRY[tool]["fn"]
    out = fn(args)
    audit_log(
        trace_id=trace_id,
        user=user,
        role=role,
        tool=tool,
        args=args,
        status="allow",
        ms=round((time.time() - t0) * 1000),
    )
    return {"ok": True, "data": out, "trace_id": trace_id}


if __name__ == "__main__":
    # 演示输出：
    # - 你会看到每次调用都打印一条审计 JSON（allow/deny 都有）
    # - 同一用户短时间连续调用会触发 rate_limit
    print("\n--- viewer 调用 inspect_schema（允许）---")
    print(guarded_call("u1", "viewer", "inspect_schema", {}))

    print("\n--- viewer 调用 export_report（拒绝：权限不足，deny 也审计）---")
    print(guarded_call("u1", "viewer", "export_report", {"format": "pdf"}))

    print("\n--- admin 调用 export_report（允许）---")
    print(guarded_call("admin_1", "admin", "export_report", {"format": "pdf"}))

    print("\n--- 同一用户快速连续调用（触发限流）---")
    for _ in range(5):
        print(guarded_call("u2", "viewer", "inspect_schema", {}))

import os
import logging

from langchain.agents import create_agent
from langchain.agents.middleware.human_in_the_loop import HumanInTheLoopMiddleware
from langchain_core.tools import tool
from langchain_community.chat_models import ChatTongyi
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command

# 配置日志
logging.basicConfig(level=logging.INFO)

# ============================
# 定义高危工具：数据导出
# ============================
@tool
def export_salary_report(department: str, format: str = "csv") -> str:
    """导出部门薪资报告（高危工具）。
    
    Args:
        department: 部门名称，如 "IT"、"Sales"
        format: 导出格式，支持 "csv" 或 "excel"
    
    Returns:
        导出结果描述
    """
    # 模拟导出操作（真实场景会访问数据库并生成文件）
    print(f"[EXPORT] 正在导出 {department} 部门的薪资报告，格式：{format}")
    return f"已成功导出 {department} 部门薪资报告（格式：{format}），包含敏感字段：姓名、工号、薪资、手机号"


@tool
def query_employee_info(employee_id: str) -> str:
    """查询员工基本信息（普通工具，不涉及敏感数据）。
    
    Args:
        employee_id: 员工工号
    
    Returns:
        员工基本信息
    """
    print(f"[QUERY] 查询员工信息：{employee_id}")
    return f"员工 {employee_id}：张三，IT部门，入职时间 2020-01-15"


# ============================
# 初始化 LLM 和 Agent
# ============================
api_key = os.getenv("DASHSCOPE_API_KEY")
if not api_key:
    raise ValueError("请先在环境变量中设置 DASHSCOPE_API_KEY")

llm = ChatTongyi(model_name="qwen-turbo", dashscope_api_key=api_key, temperature=0)

# 配置人类介入中间件：只对 export_salary_report 工具触发审批
hitl = HumanInTheLoopMiddleware(
    interrupt_on={
        "export_salary_report": True,  # 导出薪资报告需要人工确认
    }
)

# 创建 Agent
agent = create_agent(
    llm,
    tools=[export_salary_report, query_employee_info],
    middleware=(hitl,),  # 启用人类介入中间件
    system_prompt=(
        "你是一个数据分析助手。"
        "当用户需要导出薪资报告时，必须调用 export_salary_report 工具。"
        "当用户需要查询员工信息时，调用 query_employee_info 工具。"
    ),
    debug=True,
    checkpointer=MemorySaver(),  # 必须配置 checkpointer 才能支持中断后继续
)

# ============================
# 执行流程
# ============================
run_config = {
    "configurable": {"thread_id": "salary_export_demo"},
    "recursion_limit": 100,
}

# 用户请求：导出 IT 部门薪资报告（会触发人工确认）
state = agent.invoke(
    {"messages": [("human", "帮我导出 IT 部门的薪资报告，格式用 CSV")]},
    config=run_config
)

# 检查是否触发了人工审批
if "__interrupt__" in state:
    req = state["__interrupt__"][0].value
    action_req = req["action_requests"][0]
    
    print("\n" + "="*60)
    print("🚨 检测到高风险操作，需要人工确认")
    print("="*60)
    print(f"工具名称: {action_req['name']}")
    print(f"工具参数: {action_req['args']}")
    print(f"工具描述: {action_req.get('description', 'N/A')}")
    print("="*60)
    
    # 人工决策
    decision = input("\n请选择操作 (approve/reject/edit): ").strip().lower()
    
    if decision == "approve":
        print("\n✅ 已批准，正在执行工具调用...")
        cmd = Command(resume={"decisions": [{"type": "approve"}]})
    elif decision == "reject":
        print("\n❌ 已拒绝，工具不会执行")
        cmd = Command(resume={"decisions": [{"type": "reject", "message": "人工拒绝执行"}]})
    elif decision == "edit":
        # 示例：修改导出格式为 excel
        print("\n✏️ 已编辑，修改导出格式为 excel")
        edited_action = {
            "name": action_req["name"],
            "args": {"department": action_req["args"]["department"], "format": "excel"},
        }
        cmd = Command(resume={"decisions": [{"type": "edit", "edited_action": edited_action}]})
    else:
        print("\n⚠️ 无效决策，默认拒绝")
        cmd = Command(resume={"decisions": [{"type": "reject", "message": "未给出有效决策"}]})
    
    # 继续执行
    state = agent.invoke(cmd, config=run_config)

# 输出最终结果
print("\n" + "="*60)
print("📊 最终结果")
print("="*60)
print(state["messages"][-1].content)
print("="*60)

# ============================
# 验证要点
# ============================
print("\n✅ 验证要点：")
print("1. 调用 export_salary_report 时会触发人工确认")
print("2. 选择 approve 后工具会执行，可以看到 [EXPORT] 日志")
print("3. 选择 reject 后工具不会执行，不会看到 [EXPORT] 日志")
print("4. 选择 edit 后可以修改工具参数（如导出格式）")
print("5. 调用 query_employee_info 不会触发确认，直接执行")
print("\n💡 扩展建议：")
print("- 可以在 interrupt_on 中添加更多高危工具（如 delete_data、update_schema）")
print("- 可以结合审计日志记录所有确认/拒绝操作")
print("- 可以配置不同角色的审批权限（如普通用户只能查询，管理员才能导出）")

import os
import time
import logging
from typing import Dict, List

from langchain.agents import create_agent
from langchain.agents.middleware.types import AgentMiddleware, ToolCallRequest
from langchain_core.tools import tool
from langchain_community.chat_models import ChatTongyi
from langgraph.checkpoint.memory import MemorySaver

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)

# ============================
# 自定义 Middleware：日志记录
# ============================
class ToolLoggingMiddleware(AgentMiddleware):
    """记录每次工具调用的详细信息（工具名、参数、结果）"""
    
    def __init__(self, logger_name: str = "mw.tools", include_args: bool = True) -> None:
        super().__init__()
        self.logger = logging.getLogger(logger_name)
        self.include_args = include_args
        self.call_history: List[Dict] = []  # 保存所有调用历史
    
    def wrap_tool_call(self, request: ToolCallRequest, handler):
        tool_call = request.tool_call or {}
        name = tool_call.get("name")
        args = tool_call.get("args")
        
        # 记录工具调用开始
        prefix = "[ToolLoggingMiddleware]"
        if self.include_args:
            print(f"{prefix} 🔧 tool_call name={name} args={args}")
            self.logger.info("tool_call name=%s args=%s", name, args)
        else:
            print(f"{prefix} 🔧 tool_call name={name}")
            self.logger.info("tool_call name=%s", name)
        
        # 执行工具
        start_time = time.time()
        try:
            result = handler(request)
            success = True
            error_msg = None
        except Exception as e:
            result = None
            success = False
            error_msg = str(e)
            print(f"{prefix} ❌ tool_error name={name} error={error_msg}")
            self.logger.error("tool_error name=%s error=%s", name, error_msg)
            raise  # 重新抛出异常
        
        duration_ms = int((time.time() - start_time) * 1000)
        
        # 记录工具调用结果
        print(f"{prefix} ✅ tool_result name={name} type={type(result).__name__} duration={duration_ms}ms")
        self.logger.info("tool_result name=%s type=%s duration=%sms", name, type(result).__name__, duration_ms)
        
        # 保存到历史记录
        self.call_history.append({
            "tool_name": name,
            "args": args if self.include_args else None,
            "success": success,
            "error_msg": error_msg,
            "duration_ms": duration_ms,
            "timestamp": time.time()
        })
        
        return result
    
    def get_call_history(self) -> List[Dict]:
        """获取所有工具调用历史"""
        return self.call_history


# ============================
# 自定义 Middleware：性能监控
# ============================
class ToolTimingMiddleware(AgentMiddleware):
    """记录每次工具调用的耗时，识别性能瓶颈"""
    
    def __init__(self, logger_name: str = "mw.timing", threshold_ms: int = 100) -> None:
        super().__init__()
        self.logger = logging.getLogger(logger_name)
        self.threshold_ms = threshold_ms
        self.timing_stats: Dict[str, List[int]] = {}  # 保存每个工具的耗时统计
    
    def wrap_tool_call(self, request: ToolCallRequest, handler):
        tool_call = request.tool_call or {}
        name = tool_call.get("name")
        
        # 记录开始时间
        start = time.perf_counter()
        result = handler(request)
        cost_ms = int((time.perf_counter() - start) * 1000)
        
        # 保存耗时统计
        if name not in self.timing_stats:
            self.timing_stats[name] = []
        self.timing_stats[name].append(cost_ms)
        
        # 只记录超过阈值的调用
        prefix = "[ToolTimingMiddleware]"
        if cost_ms >= self.threshold_ms:
            print(f"{prefix} ⚡ tool_timing name={name} cost_ms={cost_ms} (超过阈值 {self.threshold_ms}ms)")
            self.logger.warning("tool_timing name=%s cost_ms=%s threshold=%s", name, cost_ms, self.threshold_ms)
        else:
            print(f"{prefix} ⚡ tool_timing name={name} cost_ms={cost_ms}")
            self.logger.info("tool_timing name=%s cost_ms=%s", name, cost_ms)
        
        return result
    
    def get_performance_summary(self) -> Dict[str, Dict]:
        """获取性能统计摘要"""
        summary = {}
        for tool_name, timings in self.timing_stats.items():
            if timings:
                summary[tool_name] = {
                    "total_calls": len(timings),
                    "avg_ms": sum(timings) / len(timings),
                    "min_ms": min(timings),
                    "max_ms": max(timings),
                    "total_ms": sum(timings)
                }
        return summary


# ============================
# 定义数据分析工具
# ============================
@tool
def check_schema(table_name: str) -> str:
    """检查数据表结构（模拟工具）"""
    time.sleep(0.1)  # 模拟耗时
    return f"表 {table_name} 包含字段：id, name, department, salary, hire_date"


@tool
def calculate_metrics(metric_name: str, group_by: str = None) -> str:
    """计算指标（模拟工具）"""
    time.sleep(0.2)  # 模拟耗时
    if group_by:
        return f"已计算 {metric_name}，按 {group_by} 分组：IT部门平均值=15000, Sales部门平均值=12000"
    return f"已计算 {metric_name}：总体平均值=13500"


@tool
def generate_chart(chart_type: str, data_source: str) -> str:
    """生成图表（模拟工具）"""
    time.sleep(0.15)  # 模拟耗时
    return f"已生成 {chart_type} 图表，数据来源：{data_source}，保存路径：/tmp/chart_{chart_type}.png"


# ============================
# 初始化 LLM 和 Agent
# ============================
api_key = os.getenv("DASHSCOPE_API_KEY")
if not api_key:
    raise ValueError("请先在环境变量中设置 DASHSCOPE_API_KEY")

llm = ChatTongyi(model_name="qwen-turbo", dashscope_api_key=api_key, temperature=0)

# 创建 middleware 实例
tool_logger = ToolLoggingMiddleware(logger_name="mw.tools", include_args=True)
tool_timer = ToolTimingMiddleware(logger_name="mw.timing", threshold_ms=100)

# 创建 Agent（middleware 链：日志 + 性能监控）
agent = create_agent(
    llm,
    tools=[check_schema, calculate_metrics, generate_chart],
    middleware=(tool_logger, tool_timer),  # 组合多个 middleware
    system_prompt=(
        "你是一个数据分析助手。"
        "当用户需要分析数据时，先检查表结构，然后计算指标，最后生成图表。"
    ),
    debug=True,
    checkpointer=MemorySaver(),
)

# ============================
# 执行分析任务
# ============================
run_config = {
    "configurable": {"thread_id": "data_analysis_demo"},
    "recursion_limit": 100,
}

print("="*60)
print("开始执行数据分析任务")
print("="*60)

state = agent.invoke(
    {"messages": [("human", "分析员工薪资数据：先检查 employees 表结构，然后计算各部门平均薪资，最后生成柱状图")]},
    config=run_config
)

print("\n" + "="*60)
print("📊 最终结果")
print("="*60)
print(state["messages"][-1].content)

# ============================
# 输出可观测性统计
# ============================
print("\n" + "="*60)
print("📈 工具调用历史（来自 ToolLoggingMiddleware）")
print("="*60)
import json
print(json.dumps(tool_logger.get_call_history(), ensure_ascii=False, indent=2))

print("\n" + "="*60)
print("⚡ 性能统计摘要（来自 ToolTimingMiddleware）")
print("="*60)
perf_summary = tool_timer.get_performance_summary()
print(json.dumps(perf_summary, ensure_ascii=False, indent=2))

# 识别性能瓶颈
if perf_summary:
    bottleneck = max(perf_summary.items(), key=lambda x: x[1]["avg_ms"])
    print(f"\n🔍 性能瓶颈：{bottleneck[0]} (平均耗时 {bottleneck[1]['avg_ms']:.1f}ms)")

# ============================
# 验证要点
# ============================
print("\n" + "="*60)
print("✅ 验证要点")
print("="*60)
print("1. ToolLoggingMiddleware 记录了所有工具调用的详细信息（名称、参数、结果、耗时）")
print("2. ToolTimingMiddleware 记录了每次调用的耗时，并识别超过阈值的慢调用")
print("3. 可以通过 get_call_history() 获取完整调用历史，用于错误回放")
print("4. 可以通过 get_performance_summary() 获取性能统计，识别瓶颈")
print("5. Middleware 是可组合的，可以同时启用多个（日志 + 性能 + 人工审批）")
print("\n💡 扩展建议：")
print("- 可以将 call_history 保存到数据库（如 Elasticsearch）用于全链路追踪")
print("- 可以添加错误分类逻辑（数据错误/逻辑错误/系统错误）")
print("- 可以集成 APM 平台（如 Datadog、Prometheus）进行实时监控")
print("- 可以添加告警机制（如耗时超过 P95 基线时发送通知）")

# ============================
# 目的：前端渲染层“按 JSON Schema 自适应渲染”
#
# 设计点：
# - Agent 输出统一 JSON（answer/table/bar/line/scatter...）
# - UI 层不关心“是哪个工具做的”，只按字段渲染
# - 这样可以做到：新增一个图表类型，只要新增一个字段即可
# ============================

def render_response(response_dict: dict):
    """根据 JSON Schema 字段自适应选择渲染组件"""
    if "answer" in response_dict:
        st.success("**回答：**")
        st.write(response_dict["answer"])
    if "table" in response_dict:
        df = pd.DataFrame(
            response_dict["table"]["data"],
            columns=response_dict["table"]["columns"]
        )
        st.dataframe(df, use_container_width=True)
    for chart_type in ["bar", "line", "scatter"]:
        if chart_type in response_dict:
            create_chart(response_dict[chart_type], chart_type)