""" 数据分析 Agent：集成 SQL 查询和 Python 代码执行能力 让 LLM 自动决定使用哪些工具来回答用户问题 安全特性： - SQL 注入防护：关键词黑名单 + 只读权限检查 - 恶意代码防护：危险模块黑名单 + 安全沙箱 - 数据泄露防护：查询行数限制 + 敏感字段脱敏 - 速率限制：防止恶意刷接口 - 日志审计：记录所有操作到文件 """ import os import re import json import logging import time from datetime import datetime from io import StringIO from contextlib import redirect_stdout from typing import Optional, Dict, Any from collections import defaultdict from dotenv import load_dotenv from langchain_openai import ChatOpenAI from langchain_core.prompts import ChatPromptTemplate from langchain_classic.agents import create_tool_calling_agent, AgentExecutor from langchain_core.tools import tool from langchain_community.utilities import SQLDatabase # ============================================================ # 配置日志审计系统 # ============================================================ logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', handlers=[ logging.StreamHandler(), logging.FileHandler('agent_audit.log', encoding='utf-8') ] ) log = logging.getLogger(__name__) # 加载环境变量配置文件 load_dotenv() # ============================================================ # 安全配置常量 # ============================================================ # 最大查询返回行数 MAX_QUERY_ROWS = 1000 # 查询超时时间（秒） QUERY_TIMEOUT = 30 # 速率限制配置 RATE_LIMIT_WINDOW = 60 # 时间窗口（秒） RATE_LIMIT_MAX_CALLS = 20 # 窗口内最大调用次数 # SQL 危险关键词 SQL_FORBIDDEN_KEYWORDS = [ 'INSERT', 'UPDATE', 'DELETE', 'DROP', 'CREATE', 'ALTER', 'TRUNCATE', 'GRANT', 'REVOKE', 'EXEC', 'EXECUTE' ] # Python 危险模块 PYTHON_FORBIDDEN_MODULES = [ 'os', 'subprocess', 'sys', 'shutil', 'socket', 'pickle', 'marshal', 'ctypes', 'multiprocessing' ] # ============================================================ # 速率限制器 # ============================================================ class RateLimiter: """简单的速率限制器，基于内存存储""" def __init__(self, window_seconds: int = 60, max_calls: int = 20): self.window_seconds = window_seconds self.max_calls = max_calls self.calls: Dict[str, list] = defaultdict(list) def is_allowed(self, client_id: str = 'default') -> tuple: """ 检查是否允许调用 返回: (is_allowed: bool, remaining: int, reset_time: int) """ now = time.time() calls = self.calls[client_id] # 清理过期的调用记录 calls[:] = [t for t in calls if now - t < self.window_seconds] if len(calls) >= self.max_calls: reset_time = int(self.window_seconds - (now - calls[0])) return False, 0, reset_time calls.append(now) remaining = self.max_calls - len(calls) return True, remaining, self.window_seconds # 全局速率限制器实例 rate_limiter = RateLimiter(RATE_LIMIT_WINDOW, RATE_LIMIT_MAX_CALLS) # ============================================================ # 审计日志记录器 # ============================================================ class AuditLogger: """审计日志记录器，记录所有 Agent 操作""" @staticmethod def log_operation( operation_type: str, input_data: str, output_data: str = None, success: bool = True, error: str = None, user_id: str = 'anonymous', **extra ): """记录操作日志""" log_entry = { 'timestamp': datetime.now().isoformat(), 'operation': operation_type, 'user_id': user_id, 'input': input_data[:500] if input_data else None, 'output': output_data[:500] if output_data else None, 'success': success, 'error': error, **extra } if success: log.info(f"[AUDIT] {json.dumps(log_entry, ensure_ascii=False)}") else: log.error(f"[AUDIT] {json.dumps(log_entry, ensure_ascii=False)}") audit_logger = AuditLogger() # ============================================================ # 数据库连接初始化 # ============================================================ db_uri = os.getenv("DATABASE_URI") if not db_uri: raise ValueError("DATABASE_URI 环境变量未设置，请检查 .env 配置文件") db = SQLDatabase.from_uri(db_uri) # ============================================================ # 工具 1: SQL 查询工具 # ============================================================ @tool def run_sql_query(query: str) -> str: """ 执行只读 SQL 查询，用于从数据库查询数据。 该工具只能执行 SELECT 查询，禁止执行 INSERT、UPDATE、DELETE 等修改操作。 安全措施： - 禁止危险 SQL 关键词 - 强制限制返回行数 - 查询超时控制 参数: query: SQL 查询语句（仅限 SELECT 语句） 返回: 查询结果的字符串格式 """ start_time = time.time() # 安全检查 1: SQL 关键词黑名单检测（使用单词边界匹配） query_upper = query.strip().upper() for keyword in SQL_FORBIDDEN_KEYWORDS: if re.search(rf'\b{keyword}\b', query_upper): error_msg = f"安全限制：禁止执行 {keyword} 操作，仅允许 SELECT 查询" audit_logger.log_operation( operation_type='SQL_QUERY', input_data=query, success=False, error=error_msg ) return f"❌ {error_msg}" # 安全检查 2: 强制添加 LIMIT（如果用户没有指定） if 'LIMIT' not in query_upper: query = query.rstrip(';').strip() query = f"{query} LIMIT {MAX_QUERY_ROWS}" # 安全检查 3: 验证 LIMIT 值不超过最大值 limit_match = re.search(r'LIMIT\s+(\d+)', query, re.IGNORECASE) if limit_match: limit_value = int(limit_match.group(1)) if limit_value > MAX_QUERY_ROWS: query = re.sub(r'LIMIT\s+\d+', f'LIMIT {MAX_QUERY_ROWS}', query, flags=re.IGNORECASE) try: result = db.run(query) execution_time = time.time() - start_time # 记录审计日志 audit_logger.log_operation( operation_type='SQL_QUERY', input_data=query, output_data=result[:200] if result else None, success=True, execution_time=f"{execution_time:.3f}s" ) log.info(f"SQL 查询执行成功 | 耗时: {execution_time:.3f}s") # 检查是否超时 if execution_time > QUERY_TIMEOUT: log.warning(f"查询执行时间超过阈值: {execution_time:.3f}s > {QUERY_TIMEOUT}s") return result except Exception as e: error_msg = str(e) audit_logger.log_operation( operation_type='SQL_QUERY', input_data=query, success=False, error=error_msg ) log.error(f"SQL 查询执行失败: {error_msg}") return f"❌ SQL 查询执行失败: {error_msg}" # ============================================================ # 工具 2: Python 代码执行工具 # ============================================================ @tool def execute_python_code(code: str) -> str: """ 执行 Python 代码进行数据分析、可视化或复杂计算。 【重要】已预导入以下模块，直接使用变量名即可，无需 import： - plt: matplotlib.pyplot（绑图库） - pd: pandas（数据分析库） - np: numpy（数值计算库） - json: json（JSON处理） 示例用法： # 直接使用 plt，不要写 import matplotlib.pyplot as plt plt.figure(figsize=(10, 6)) plt.bar(['A', 'B'], [10, 20]) plt.savefig('chart.png') # 直接使用 pd 和 np df = pd.DataFrame({'name': ['张三', '李四'], 'age': [25, 30]}) arr = np.array([1, 2, 3, 4, 5]) 安全限制： - 禁止使用 import 语句（已预导入常用库） - 禁止使用 eval()、exec()、open() 等危险函数 参数: code: Python 代码字符串（直接使用预导入的 plt/pd/np/json） 返回: 代码执行的输出结果 """ start_time = time.time() # 安全检查 1: 禁止危险模块 for module in PYTHON_FORBIDDEN_MODULES: if re.search(rf'\bimport\s+{module}\b', code) or re.search(rf'\bfrom\s+{module}\b', code): error_msg = f"安全限制：禁止导入 {module} 模块" audit_logger.log_operation( operation_type='PYTHON_CODE', input_data=code[:200], success=False, error=error_msg ) return f"❌ {error_msg}" # 安全检查 2: 禁止危险内置函数 dangerous_functions = ['eval', 'exec', 'compile', 'open', '__import__'] for func in dangerous_functions: if re.search(rf'\b{func}\s*\(', code): error_msg = f"安全限制：禁止使用 {func}() 函数" audit_logger.log_operation( operation_type='PYTHON_CODE', input_data=code[:200], success=False, error=error_msg ) return f"❌ {error_msg}" # 配置 matplotlib 中文字体 import matplotlib import matplotlib.pyplot as plt chinese_fonts = ['SimHei', 'Microsoft YaHei', 'STSong', 'SimSun', 'KaiTi', 'FangSong'] available_fonts = [f.name for f in matplotlib.font_manager.fontManager.ttflist] for font in chinese_fonts: if font in available_fonts: plt.rcParams['font.sans-serif'] = [font] break plt.rcParams['axes.unicode_minus'] = False # 创建受限的执行环境（白名单模式） safe_builtins = { # 基础函数 'print': print, 'len': len, 'range': range, 'enumerate': enumerate, 'zip': zip, 'map': map, 'filter': filter, 'sorted': sorted, 'reversed': reversed, # 类型 'list': list, 'dict': dict, 'tuple': tuple, 'set': set, 'str': str, 'int': int, 'float': float, 'bool': bool, 'bytes': bytes, # 数学函数 'sum': sum, 'min': min, 'max': max, 'abs': abs, 'round': round, 'pow': pow, 'divmod': divmod, # 类型检查 'isinstance': isinstance, 'type': type, 'hasattr': hasattr, 'getattr': getattr, 'callable': callable, # 异常 'Exception': Exception, 'ValueError': ValueError, 'TypeError': TypeError, 'KeyError': KeyError, 'IndexError': IndexError, 'AttributeError': AttributeError, 'RuntimeError': RuntimeError, 'StopIteration': StopIteration, # 其他安全函数 'repr': repr, 'hash': hash, 'id': id, } # 允许的安全模块白名单（用于受限的 __import__） ALLOWED_MODULES = {'pandas', 'numpy', 'matplotlib', 'matplotlib.pyplot', 'json'} def safe_import(name, *args, **kwargs): """安全的 import 函数，只允许导入白名单中的模块""" if name.split('.')[0] not in {'pandas', 'numpy', 'matplotlib', 'json'}: raise ImportError(f"安全限制：禁止导入模块 '{name}'，请直接使用预导入的 plt/pd/np/json") return __import__(name, *args, **kwargs) safe_builtins['__import__'] = safe_import local_vars = { 'pd': __import__('pandas'), 'np': __import__('numpy'), 'plt': plt, 'json': json, } # 捕获标准输出 output_buffer = StringIO() try: with redirect_stdout(output_buffer): exec(code, {"__builtins__": safe_builtins}, local_vars) output = output_buffer.getvalue() execution_time = time.time() - start_time if not output: result_vars = {k: v for k, v in local_vars.items() if not k.startswith('_') and k not in ['pd', 'np', 'plt', 'json']} if result_vars: output = "执行成功，生成的变量:\n" for var_name, var_value in result_vars.items(): output += f" {var_name} = {repr(var_value)[:200]}\n" else: output = "✅ 代码执行成功（无输出）" # 记录审计日志 audit_logger.log_operation( operation_type='PYTHON_CODE', input_data=code[:200], output_data=output[:200], success=True, execution_time=f"{execution_time:.3f}s" ) log.info(f"Python 代码执行成功 | 耗时: {execution_time:.3f}s") return output except Exception as e: error_msg = str(e) audit_logger.log_operation( operation_type='PYTHON_CODE', input_data=code[:200], success=False, error=error_msg ) log.error(f"Python 代码执行失败: {error_msg}") return f"❌ Python 代码执行失败: {error_msg}" # ============================================================ # Agent 创建 # ============================================================ def create_data_analysis_agent(model_name: str = None, temperature: float = 0): """ 创建数据分析 Agent，集成 SQL 查询和 Python 代码执行能力 参数: model_name: 使用的模型名称，默认从环境变量读取 temperature: 模型温度，默认为 0（更确定性的输出） """ # 从环境变量读取配置 api_key = os.getenv("OPENAI_API_KEY") api_base = os.getenv("OPENAI_API_BASE") model_name = os.getenv("MODEL_NAME") if not api_key: raise ValueError("OPENAI_API_KEY 环境变量未设置，请检查 .env 配置文件") # 初始化 LLM llm_config = { "model": model_name, "temperature": temperature, } if api_base: llm_config["openai_api_base"] = api_base log.info(f"使用自定义 API 端点: {api_base}") llm = ChatOpenAI(**llm_config) log.info(f"初始化模型: {model_name}") # 定义工具列表 tools = [run_sql_query, execute_python_code] # 创建 Prompt 模板 prompt = ChatPromptTemplate.from_messages([ ("system", """你是一个专业的数据分析助手，可以访问公司数据库并执行 Python 代码进行数据分析。 数据库包含以下表： 1. employees (员工表) - id: 员工ID - name: 姓名 - department: 所属部门 - salary: 月薪 - hire_date: 入职日期 2. products (产品表) - id: 产品ID - product_name: 商品名称 - category: 商品分类 - price: 单价 - stock: 当前库存量 3. orders (订单表) - id: 订单ID - employee_id: 下单员工ID - product_id: 购买商品ID - quantity: 购买数量 - order_date: 下单日期 工作流程： 1. 理解用户的问题 2. 决定需要哪些数据 3. 使用 run_sql_query 工具查询数据库 4. 如果需要复杂分析或可视化，使用 execute_python_code 工具 5. 用自然语言总结结果 注意事项： - SQL 查询仅支持 SELECT 语句 - 使用 Python 时，已预导入 pandas as pd、numpy as np、matplotlib.pyplot as plt - 如果查询结果需要进一步分析，可以在 Python 代码中使用 SQL 查询的结果 - 始终用中文回答问题 - 提供清晰的分析结论和建议 - 注意数据安全，不要泄露敏感信息 """), ("human", "{input}"), ("placeholder", "{agent_scratchpad}"), ]) # 创建 Agent agent = create_tool_calling_agent(llm, tools, prompt) # 创建 AgentExecutor agent_executor = AgentExecutor( agent=agent, tools=tools, verbose=True, handle_parsing_errors=True, max_iterations=10, ) log.info("Agent 创建成功") return agent_executor # ============================================================ # 主函数：统一 Agent 入口（带速率限制） # ============================================================ def run_data_analysis(user_query: str, model_name: str = None, client_id: str = 'default') -> str: """ 统一数据分析 Agent 入口 参数: user_query: 用户的自然语言查询问题 model_name: 可选，指定使用的模型名称 client_id: 客户端标识，用于速率限制 返回: Agent 的分析结果 """ # 速率限制检查 is_allowed, remaining, reset_time = rate_limiter.is_allowed(client_id) if not is_allowed: error_msg = f"请求过于频繁，请在 {reset_time} 秒后重试" log.warning(f"速率限制触发: client_id={client_id}") return f"❌ {error_msg}" # 记录用户查询 audit_logger.log_operation( operation_type='USER_QUERY', input_data=user_query, client_id=client_id ) try: agent = create_data_analysis_agent(model_name=model_name) result = agent.invoke({"input": user_query}) log.info(f"查询处理完成: {user_query[:50]}... | 剩余调用次数: {remaining}") return result["output"] except Exception as e: error_msg = str(e) audit_logger.log_operation( operation_type='AGENT_ERROR', input_data=user_query, success=False, error=error_msg, client_id=client_id ) log.error(f"Agent 执行失败: {error_msg}") return f"❌ Agent 执行失败: {error_msg}" # ============================================================ # 测试入口 # ============================================================ if __name__ == "__main__": print("=" * 60) print("数据分析 Agent 测试（安全增强版）") print("=" * 60) # 测试: 绘制图表 print("\n【测试】绘制图表 - 各部门薪资分布柱状图:") result = run_data_analysis( "查询各部门的平均薪资，并用 matplotlib 绘制柱状图，" "图表标题为'各部门平均薪资分布'，x轴为部门名称，y轴为平均薪资，" "保存图表到 d:/agentlearning/lqq-agent-study/salary_chart.png" ) print(result)