langchain一个agent集成sql查询和Python代码执行能力

01_langchain_study/01_langchain_task.py

@@ -0,0 +1,593 @@
+"""
+数据分析 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("\n【直接回答版:")
+    result = run_data_analysis(
+        "今天上海天气怎么样，会来台风吗"
+    )
+    print(result)