{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "02b24f3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ DeepSeek 模型初始化完成\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from langchain_openai import ChatOpenAI\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "\n",
    "# 通过环境变量读取 API Key，避免硬编码泄露风险\n",
    "llm = ChatOpenAI(\n",
    "    model=os.getenv(\"DEEPSEEK_MODEL\"),                          # DeepSeek 的对话模型\n",
    "    api_key=os.getenv(\"DEEPSEEK_API_KEY\"),          # 从环境变量加载\n",
    "    base_url=os.getenv(\"DEEPSEEK_API_BASE\"),            # DeepSeek API 地址\n",
    "    temperature=0,                                   # 设为 0 保证 SQL 生成的确定性\n",
    ")\n",
    "\n",
    "print(\"✅ DeepSeek 模型初始化完成\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3a603f66",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ 数据库初始化完成\n",
      "   - employees 表：5 条员工记录\n",
      "   - products  表：4 条产品记录\n",
      "   - orders    表：5 条订单记录\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import mysql.connector\n",
    "\n",
    "# ============================================================\n",
    "# 第一步：建立数据库连接\n",
    "# ============================================================\n",
    "# 使用环境变量管理敏感信息，这是生产级代码的基本规范\n",
    "conn = mysql.connector.connect(\n",
    "    host=os.getenv(\"DB_HOST\", \"127.0.0.1\"),\n",
    "    port=int(os.getenv(\"DB_PORT\", 3308)),\n",
    "    user=os.getenv(\"DB_USER\", \"root\"),\n",
    "    password=os.getenv(\"DB_PASSWORD\", \"admin\"),\n",
    "    database=os.getenv(\"DB_NAME\", \"analytics_demo\"),\n",
    ")\n",
    "cursor = conn.cursor()\n",
    "\n",
    "# ============================================================\n",
    "# 第二步：创建表结构\n",
    "# ============================================================\n",
    "\n",
    "# 员工表：存储公司内部人员信息\n",
    "# 用于分析：部门人数分布、薪资结构、入职趋势等\n",
    "cursor.execute(\"\"\"\n",
    "CREATE TABLE IF NOT EXISTS employees (\n",
    "    id          INT PRIMARY KEY AUTO_INCREMENT,  -- 自增主键\n",
    "    name        VARCHAR(50)  NOT NULL,           -- 姓名\n",
    "    department  VARCHAR(50)  NOT NULL,           -- 所属部门\n",
    "    salary      DECIMAL(10,2) NOT NULL,          -- 月薪（精确到分）\n",
    "    hire_date   DATE                             -- 入职日期\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;\n",
    "\"\"\")\n",
    "\n",
    "# 产品表：存储在售商品信息\n",
    "# 用于分析：品类销售、价格分布、库存预警等\n",
    "cursor.execute(\"\"\"\n",
    "CREATE TABLE IF NOT EXISTS products (\n",
    "    id           INT PRIMARY KEY AUTO_INCREMENT,  -- 自增主键\n",
    "    product_name VARCHAR(100) NOT NULL,           -- 商品名称\n",
    "    category     VARCHAR(50)  NOT NULL,           -- 商品分类\n",
    "    price        DECIMAL(10,2) NOT NULL,          -- 单价\n",
    "    stock        INT DEFAULT 0                    -- 当前库存量\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;\n",
    "\"\"\")\n",
    "\n",
    "# 订单表：记录每笔交易\n",
    "# 用于分析：销售趋势、员工绩效、产品热度等\n",
    "cursor.execute(\"\"\"\n",
    "CREATE TABLE IF NOT EXISTS orders (\n",
    "    id           INT PRIMARY KEY AUTO_INCREMENT,  -- 自增主键\n",
    "    employee_id  INT NOT NULL,                    -- 下单员工（外键）\n",
    "    product_id   INT NOT NULL,                    -- 购买商品（外键）\n",
    "    quantity     INT NOT NULL,                    -- 购买数量\n",
    "    order_date   DATE NOT NULL,                   -- 下单日期\n",
    "    FOREIGN KEY (employee_id) REFERENCES employees(id),   -- 关联员工表\n",
    "    FOREIGN KEY (product_id)  REFERENCES products(id)     -- 关联产品表\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;\n",
    "\"\"\")\n",
    "\n",
    "# ============================================================\n",
    "# 第三步：插入演示数据\n",
    "# ============================================================\n",
    "\n",
    "# 员工数据：5 个员工，分属 3 个部门\n",
    "employees_data = [\n",
    "    (1, \"张三\", \"技术部\",   20000.00, \"2023-01-15\"),\n",
    "    (2, \"李四\", \"销售部\",   11000.00, \"2023-02-20\"),\n",
    "    (3, \"王五\", \"技术部\",   16000.00, \"2022-11-10\"),\n",
    "    (4, \"赵六\", \"人力资源\", 5000.00, \"2023-03-01\"),\n",
    "    (5, \"钱七\", \"销售部\",   17000.00, \"2022-12-05\"),\n",
    "]\n",
    "\n",
    "# 产品数据：4 款商品，覆盖 2 个品类\n",
    "products_data = [\n",
    "    (1, \"笔记本电脑\", \"电子产品\", 6999.00, 500),\n",
    "    (2, \"机械键盘\",   \"电子产品\", 399.00,  1000),\n",
    "    (3, \"办公椅\",     \"办公用品\", 499.00,  300),\n",
    "    (4, \"显示器\",     \"电子产品\", 1200.00, 400),\n",
    "]\n",
    "\n",
    "# 订单数据：5 笔订单，模拟真实购买行为\n",
    "orders_data = [\n",
    "    (1, 1, 1, 2, \"2024-01-15\"),  # 张三买了 2 台笔记本\n",
    "    (2, 2, 2, 15, \"2024-01-16\"),  # 李四买了 15 个键盘\n",
    "    (3, 3, 1, 10, \"2024-01-17\"),  # 王五买了 10 台笔记本\n",
    "    (4, 5, 3, 6, \"2024-01-18\"),  # 钱七买了 6 把办公椅\n",
    "    (5, 2, 4, 5, \"2024-01-19\"),  # 李四买了 5 台显示器\n",
    "]\n",
    "\n",
    "# executemany 批量插入，比逐条 insert 高效得多\n",
    "cursor.executemany(\"INSERT IGNORE INTO employees VALUES (%s,%s,%s,%s,%s)\", employees_data)\n",
    "cursor.executemany(\"INSERT IGNORE INTO products  VALUES (%s,%s,%s,%s,%s)\", products_data)\n",
    "cursor.executemany(\"INSERT IGNORE INTO orders     VALUES (%s,%s,%s,%s,%s)\", orders_data)\n",
    "\n",
    "conn.commit()\n",
    "conn.close()\n",
    "\n",
    "print(\"✅ 数据库初始化完成\")\n",
    "print(\"   - employees 表：5 条员工记录\")\n",
    "print(\"   - products  表：4 条产品记录\")\n",
    "print(\"   - orders    表：5 条订单记录\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e02aa540",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "数据库连接成功\n",
      "   可用表：['employees', 'orders', 'products']\n",
      "\n",
      "SQL 工具包已加载（4 个工具）：\n",
      "   - sql_db_query\n",
      "   - sql_db_schema\n",
      "   - sql_db_list_tables\n",
      "   - sql_db_query_checker\n",
      "\n",
      "NL2SQL Agent 创建完成，可以开始提问了！\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "from langchain_community.agent_toolkits import SQLDatabaseToolkit\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain.agents import create_agent  # langchain 1.3.1 的新 API\n",
    "\n",
    "# ============================================================\n",
    "# 第一步：连接数据库\n",
    "# ============================================================\n",
    "# SQLDatabase.from_uri 接受标准的数据库连接 URI\n",
    "# LangChain 内部会用 SQLAlchemy 管理连接池\n",
    "db_uri = (\n",
    "    f\"mysql+pymysql://{os.getenv('DB_USER')}:{os.getenv('DB_PASSWORD')}\"\n",
    "    f\"@{os.getenv('DB_HOST')}:{os.getenv('DB_PORT')}/{os.getenv('DB_NAME')}\"\n",
    ")\n",
    "db = SQLDatabase.from_uri(db_uri)\n",
    "\n",
    "# 验证连接：打印可用的表名\n",
    "print(f\"数据库连接成功\")\n",
    "print(f\"   可用表：{db.get_usable_table_names()}\")\n",
    "\n",
    "# ============================================================\n",
    "# 第二步：初始化大模型\n",
    "# ============================================================\n",
    "llm = ChatOpenAI(\n",
    "    model=os.getenv(\"DEEPSEEK_MODEL\"),                          # DeepSeek 的对话模型\n",
    "    api_key=os.getenv(\"DEEPSEEK_API_KEY\"),          # 从环境变量加载\n",
    "    base_url=os.getenv(\"DEEPSEEK_API_BASE\"),            # DeepSeek API 地址\n",
    "    temperature=0,                                   # 设为 0 保证 SQL 生成的确定性\n",
    ")\n",
    "\n",
    "# ============================================================\n",
    "# 第三步：创建 SQL 工具包\n",
    "# ============================================================\n",
    "# SQLDatabaseToolkit 会自动注册 4 个工具：\n",
    "#   1. sql_db_list_tables    — 列出数据库中所有表\n",
    "#   2. sql_db_schema         — 获取指定表的 DDL 结构\n",
    "#   3. sql_db_query_checker  — 检查 SQL 语法是否正确\n",
    "#   4. sql_db_query          — 执行 SQL 并返回结果\n",
    "toolkit = SQLDatabaseToolkit(db=db, llm=llm)\n",
    "tools = toolkit.get_tools()\n",
    "\n",
    "print(f\"\\nSQL 工具包已加载（{len(tools)} 个工具）：\")\n",
    "for t in tools:\n",
    "    print(f\"   - {t.name}\")\n",
    "\n",
    "# ============================================================\n",
    "# 第四步：定义 System Prompt（Agent 的\"岗位说明书\"）\n",
    "# ============================================================\n",
    "SQL_AGENT_PROMPT = \"\"\"你是一名专业的 SQL 数据分析师。\n",
    "\n",
    "## 工作流程\n",
    "1. 先用 sql_db_list_tables 查看数据库中有哪些表\n",
    "2. 用 sql_db_schema 获取相关表的字段结构和类型\n",
    "3. 生成 SQL 之前，用 sql_db_query_checker 检查语法\n",
    "4. 确认无误后，用 sql_db_query 执行查询\n",
    "5. 用中文总结查询结果，给出简洁的业务洞察\n",
    "\n",
    "## 约束\n",
    "- 只使用数据库中实际存在的表和字段，不要凭空编造\n",
    "- 单次查询结果限制在 50 条以内\n",
    "- 如果查询出错，分析错误原因后重新生成 SQL\n",
    "- 回答要简洁专业，不要啰嗦\n",
    "\"\"\"\n",
    "\n",
    "# ============================================================\n",
    "# 第五步：组装 Agent（langchain 1.3.1 一行搞定）\n",
    "# ============================================================\n",
    "# create_agent 返回一个可直接调用的 Runnable，无需再套 AgentExecutor\n",
    "# 它内部自动处理工具调用、循环迭代、错误重试等逻辑\n",
    "sql_agent = create_agent(\n",
    "    model=llm,\n",
    "    tools=tools,\n",
    "    system_prompt=SQL_AGENT_PROMPT,\n",
    ")\n",
    "\n",
    "print(\"\\nNL2SQL Agent 创建完成，可以开始提问了！\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bd07bbe0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "回答：公司共有 **5 名员工**。\n"
     ]
    }
   ],
   "source": [
    "# 简单问题：Agent 只需一条 COUNT SQL 就能搞定\n",
    "# create_agent 的 invoke 接口使用 messages 格式\n",
    "response = await sql_agent.ainvoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"公司一共有多少名员工？\"}]\n",
    "})\n",
    "\n",
    "# 最终回答在最后一条消息里\n",
    "final_msg = response[\"messages\"][-1]\n",
    "print(f\"回答：{final_msg.content}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "55db4cad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "回答：查询结果如下：\n",
      "\n",
      "| 姓名 | 薪资 |\n",
      "|------|------|\n",
      "| 张三 | 20,000.00 元 |\n",
      "| 王五 | 16,000.00 元 |\n",
      "\n",
      "**业务洞察：**\n",
      "- 技术部共有 **2 名员工**，平均薪资为 **18,000 元**。\n",
      "- 张三（20,000 元）薪资最高，比王五（16,000 元）高出 **25%**。\n"
     ]
    }
   ],
   "source": [
    "# 中等难度：需要 JOIN + WHERE + GROUP BY\n",
    "response = await sql_agent.ainvoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"列出技术部所有员工的姓名和薪资，按薪资从高到低排序\"}]\n",
    "})\n",
    "final_msg = response[\"messages\"][-1]\n",
    "print(f\"回答：{final_msg.content}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f884cdc0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "回答：## 查询结果\n",
      "\n",
      "**销售部员工**共有2人：**李四** 和 **钱七**\n",
      "\n",
      "### 汇总数据\n",
      "| 指标 | 数值 |\n",
      "|------|------|\n",
      "| 总订单数 | **3 笔** |\n",
      "| 订单总金额 | **¥14,979.00** |\n",
      "\n",
      "### 员工明细\n",
      "| 员工 | 订单数 | 订单金额 |\n",
      "|------|:------:|:--------:|\n",
      "| 李四 | 2 笔 | ¥11,985.00 |\n",
      "| 钱七 | 1 笔 | ¥2,994.00 |\n",
      "\n",
      "### 业务洞察\n",
      "- 销售部两位员工共下3笔订单，总金额约 **1.5万元**\n",
      "- **李四**贡献了约 **80%** 的订单金额，是销售主力\n",
      "- 建议可进一步了解钱七的客户资源情况，提升其订单量\n"
     ]
    }
   ],
   "source": [
    "# 高难度：需要 JOIN 三张表 + 聚合计算\n",
    "response = await sql_agent.ainvoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"销售部的员工总共下了多少订单？订单总金额是多少？\"}]\n",
    "})\n",
    "final_msg = response[\"messages\"][-1]\n",
    "print(f\"回答：{final_msg.content}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a9b0e412",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "最终回答：每个部门各有多少人？\n",
      "\n",
      "步骤 2 [AIMessage]：\n",
      "  调用工具：sql_db_list_tables\n",
      "  参数：{}\n",
      "  工具返回：employees, orders, products\n",
      "\n",
      "步骤 4 [AIMessage]：\n",
      "  调用工具：sql_db_schema\n",
      "  参数：{'table_names': 'employees'}\n",
      "  工具返回：\n",
      "CREATE TABLE employees (\n",
      "\tid INTEGER NOT NULL AUTO_INCREMENT, \n",
      "\tname VARCHAR(50) NOT NULL, \n",
      "\tdepartment VARCHAR(50) NOT NULL, \n",
      "\tsalary DECIMAL(10, 2) NOT NULL, \n",
      "\thire_date DATE, \n",
      "\tPRIMARY KEY (id)\n",
      ")D\n",
      "\n",
      "步骤 6 [AIMessage]：\n",
      "  调用工具：sql_db_query_checker\n",
      "  参数：{'query': 'SELECT department, COUNT(*) AS 人数 FROM employees GROUP BY department ORDER BY 人数 DESC'}\n",
      "  工具返回：SELECT department, COUNT(*) AS 人数 FROM employees GROUP BY department ORDER BY 人数 DESC\n",
      "\n",
      "步骤 8 [AIMessage]：\n",
      "  调用工具：sql_db_query\n",
      "  参数：{'query': 'SELECT department, COUNT(*) AS 人数 FROM employees GROUP BY department ORDER BY 人数 DESC'}\n",
      "  工具返回：[('技术部', 2), ('销售部', 2), ('人力资源', 1)]\n",
      "\n",
      "最终回答：查询结果如下：\n",
      "\n",
      "| 部门 | 人数 |\n",
      "|------|:----:|\n",
      "| 技术部 | 2 人 |\n",
      "| 销售部 | 2 人 |\n",
      "| 人力资源 | 1 人 |\n",
      "\n",
      "**共计 5 人**，分布在 3 个部门。技术部和销售部各有 2 人，人力资源部有 1 人。\n"
     ]
    }
   ],
   "source": [
    "# 遍历所有消息，还原 Agent 的完整推理链\n",
    "response = await sql_agent.ainvoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"每个部门各有多少人？\"}]\n",
    "})\n",
    "\n",
    "for i, msg in enumerate(response[\"messages\"]):\n",
    "    msg_type = msg.__class__.__name__\n",
    "\n",
    "    # AIMessage 中如果有 tool_calls，说明 Agent 调用了工具\n",
    "    if hasattr(msg, \"tool_calls\") and msg.tool_calls:\n",
    "        print(f\"\\n步骤 {i+1} [{msg_type}]：\")\n",
    "        for tc in msg.tool_calls:\n",
    "            print(f\"  调用工具：{tc['name']}\")\n",
    "            print(f\"  参数：{tc.get('args', {})}\")\n",
    "\n",
    "    # ToolMessage 是工具返回的结果\n",
    "    elif msg_type == \"ToolMessage\":\n",
    "        print(f\"  工具返回：{msg.content[:200]}\")\n",
    "\n",
    "    # AIMessage 的最终文本回答\n",
    "    elif msg.content:\n",
    "        print(f\"\\n最终回答：{msg.content}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "32ca7d76",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ 数据加载完成\n",
      "   员工表：5 行 × 5 列\n",
      "   产品表：4 行 × 5 列\n",
      "   订单表：5 行 × 5 列\n",
      "\n",
      "📋 员工表示例：\n",
      " id name department  salary  hire_date\n",
      "  1   张三        技术部 20000.0 2023-01-15\n",
      "  2   李四        销售部 11000.0 2023-02-20\n",
      "  3   王五        技术部 16000.0 2022-11-10\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "\n",
    "# ============================================================\n",
    "# 配置 matplotlib 中文显示\n",
    "# ============================================================\n",
    "# Windows 用 SimHei（黑体），macOS 用 PingFang SC\n",
    "# 如果还是乱码，试试安装 fonts-noto-cjk 并清除缓存\n",
    "plt.rcParams[\"font.sans-serif\"] = [\"SimHei\", \"PingFang SC\", \"DejaVu Sans\"]\n",
    "plt.rcParams[\"axes.unicode_minus\"] = False  # 解决负号显示为方块的问题\n",
    "\n",
    "# ============================================================\n",
    "# 从数据库加载数据到 Pandas DataFrame\n",
    "# ============================================================\n",
    "# 用 SQLAlchemy engine 复用连接，避免重复创建连接池\n",
    "from sqlalchemy import create_engine\n",
    "\n",
    "engine = create_engine(db_uri)\n",
    "\n",
    "employees_df = pd.read_sql(\"SELECT * FROM employees\", engine)\n",
    "products_df  = pd.read_sql(\"SELECT * FROM products\", engine)\n",
    "orders_df    = pd.read_sql(\"SELECT * FROM orders\", engine)\n",
    "\n",
    "print(\"✅ 数据加载完成\")\n",
    "print(f\"   员工表：{len(employees_df)} 行 × {len(employees_df.columns)} 列\")\n",
    "print(f\"   产品表：{len(products_df)} 行 × {len(products_df.columns)} 列\")\n",
    "print(f\"   订单表：{len(orders_df)} 行 × {len(orders_df.columns)} 列\")\n",
    "print(f\"\\n📋 员工表示例：\")\n",
    "print(employees_df.head(3).to_string(index=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "01276a4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Python 代码执行沙箱创建成功\n"
     ]
    }
   ],
   "source": [
    "import traceback\n",
    "from io import StringIO\n",
    "from contextlib import redirect_stdout\n",
    "from langchain.tools import tool\n",
    "\n",
    "# ============================================================\n",
    "# 定义沙箱的\"白名单\"——只有这些库和数据可以被代码访问\n",
    "# ============================================================\n",
    "# 这是一种简单的安全策略：不在白名单里的东西，代码碰不到\n",
    "SANDBOX_GLOBALS = {\n",
    "    # 数据：Agent 可以分析这三张表\n",
    "    \"employees_df\": employees_df,\n",
    "    \"products_df\":  products_df,\n",
    "    \"orders_df\":    orders_df,\n",
    "    # 工具库：Agent 可以用这些库做分析和画图\n",
    "    \"pd\":  pd,       # pandas — 数据处理\n",
    "    \"plt\": plt,      # matplotlib — 基础绑图\n",
    "    \"sns\": sns,      # seaborn — 统计可视化\n",
    "    \"np\":  np,       # numpy — 数值计算\n",
    "}\n",
    "\n",
    "\n",
    "@tool\n",
    "def execute_python_code(code: str) -> str:\n",
    "    \"\"\"\n",
    "    执行 Python 代码进行数据分析和可视化。\n",
    "\n",
    "    可用变量：\n",
    "      - employees_df: 员工表 DataFrame（字段：id, name, department, salary, hire_date）\n",
    "      - products_df:  产品表 DataFrame（字段：id, product_name, category, price, stock）\n",
    "      - orders_df:    订单表 DataFrame（字段：id, employee_id, product_id, quantity, order_date）\n",
    "      - pd:   pandas 库\n",
    "      - plt:  matplotlib.pyplot\n",
    "      - sns:  seaborn\n",
    "      - np:   numpy\n",
    "\n",
    "    使用示例：\n",
    "      # 统计各部门平均薪资\n",
    "      result = employees_df.groupby('department')['salary'].mean()\n",
    "      print(result)\n",
    "\n",
    "      # 画柱状图\n",
    "      plt.figure(figsize=(10, 6))\n",
    "      employees_df.groupby('department')['salary'].mean().plot(kind='bar')\n",
    "      plt.title('各部门平均薪资')\n",
    "      plt.show()\n",
    "    \"\"\"\n",
    "    # 准备隔离的执行环境\n",
    "    # globals_dict 提供白名单变量，locals_dict 收集执行过程中产生的新变量\n",
    "    exec_globals = dict(SANDBOX_GLOBALS)\n",
    "    exec_locals = {}\n",
    "\n",
    "    # 用 StringIO 捕获 print() 的输出\n",
    "    # 这样 Agent 生成的代码里所有的 print 语句都会被收集\n",
    "    output_buffer = StringIO()\n",
    "\n",
    "    try:\n",
    "        # redirect_stdout 会把标准输出重定向到我们的 buffer\n",
    "        with redirect_stdout(output_buffer):\n",
    "            exec(code, exec_globals, exec_locals)\n",
    "\n",
    "        result = output_buffer.getvalue()\n",
    "\n",
    "        # 如果代码没有 print 任何东西，给个默认提示\n",
    "        if not result.strip():\n",
    "            result = \"✅ 代码执行成功（无文本输出，可能已生成图表）\"\n",
    "\n",
    "        return f\"执行成功：\\n{result}\"\n",
    "\n",
    "    except Exception as e:\n",
    "        # 出错时返回完整的错误堆栈，方便 Agent 自我修正\n",
    "        error_detail = traceback.format_exc()\n",
    "        return f\"❌ 执行出错：{e}\\n\\n{error_detail}\"\n",
    "\n",
    "\n",
    "print(\"✅ Python 代码执行沙箱创建成功\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e8e2f648",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "数据可视化 Agent 创建完成\n"
     ]
    }
   ],
   "source": [
    "from langchain.agents import create_agent\n",
    "\n",
    "# ============================================================\n",
    "# 定义可视化 Agent 的 System Prompt\n",
    "# ============================================================\n",
    "VISUALIZATION_PROMPT = \"\"\"你是一名资深数据分析师，精通 Python、Pandas 和 Matplotlib 数据可视化。\n",
    "\n",
    "## 可用数据\n",
    "1. employees_df — 员工表（字段：id, name, department, salary, hire_date）\n",
    "2. products_df  — 产品表（字段：id, product_name, category, price, stock）\n",
    "3. orders_df    — 订单表（字段：id, employee_id, product_id, quantity, order_date）\n",
    "\n",
    "## 工作流程\n",
    "1. 理解用户的分析需求\n",
    "2. 用 execute_python_code 工具编写并执行 Python 代码\n",
    "3. 先做数据探索（head、describe、info），再做深入分析\n",
    "4. 用中文解释分析结果，给出业务洞察\n",
    "\n",
    "## 代码规范\n",
    "- 绑图前设置中文字体：plt.rcParams['font.sans-serif'] = ['SimHei', 'PingFang SC', 'DejaVu Sans']\n",
    "- 设置 plt.rcParams['axes.unicode_minus'] = False\n",
    "- 图表尺寸统一用 plt.figure(figsize=(10, 6))\n",
    "- 必须添加标题、坐标轴标签，让图表自解释\n",
    "- 用 print() 输出关键统计量，不要只画图不说话\n",
    "- 图表标题用英文（避免渲染问题），但用中文向用户解释结果\n",
    "\n",
    "## 注意事项\n",
    "- 每次只执行一段完整的代码，不要拆成多段\n",
    "- 先探索数据结构，再做分析——不要上来就画图\n",
    "- 结果要有业务洞察，不只是\"最大值是 XXX\"\n",
    "\"\"\"\n",
    "\n",
    "# ============================================================\n",
    "# 组装可视化 Agent（langchain 1.3.1 写法）\n",
    "# ============================================================\n",
    "viz_tools = [execute_python_code]\n",
    "\n",
    "# 同样用 create_agent 一行搞定，和 SQL Agent 的创建方式完全一致\n",
    "visualization_agent = create_agent(\n",
    "    model=llm,\n",
    "    tools=viz_tools,\n",
    "    system_prompt=VISUALIZATION_PROMPT,\n",
    ")\n",
    "\n",
    "print(\"数据可视化 Agent 创建完成\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "59ebc21f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "findfont: Failed to find font weight bold, now using 400.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 36948 characters]",
      "text/plain": [
       "<Figure size 1400x1000 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 117756 characters]",
      "text/plain": [
       "<Figure size 1400x1000 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "分析结果：\n",
      "---\n",
      "\n",
      "## 📊 员工薪资分布分析报告\n",
      "\n",
      "### 一、数据概况\n",
      "\n",
      "公司共有 **5 名员工**，分布在 **3 个部门**（技术部、销售部、人力资源），薪资数据完整无缺失。\n",
      "\n",
      "### 二、核心统计指标\n",
      "\n",
      "| 统计量 | 数值 | 说明 |\n",
      "|--------|------|------|\n",
      "| **均值（Mean）** | **¥13,800** | 员工平均薪资水平 |\n",
      "| **中位数（Median）** | **¥16,000** | 中间位置的薪资水平 |\n",
      "| **最小值（Min）** | **¥5,000** | 赵六（人力资源部） |\n",
      "| **最大值（Max）** | **¥20,000** | 张三（技术部） |\n",
      "| **极差（Range）** | **¥15,000** | 薪资差距非常大 |\n",
      "| **标准差（Std）** | **¥5,890.67** | 薪资离散程度较高 |\n",
      "| **Q1（25%分位）** | **¥11,000** | 下四分位薪资 |\n",
      "| **Q3（75%分位）** | **¥17,000** | 上四分位薪资 |\n",
      "| **偏度（Skewness）** | **-0.86** | 左偏态分布 |\n",
      "| **峰度（Kurtosis）** | **-0.04** | 接近正态分布 |\n",
      "\n",
      "### 三、关键洞察\n",
      "\n",
      "#### 1️⃣ 薪资呈**左偏分布**\n",
      "> **均值（¥13,800） < 中位数（¥16,000）**，说明大部分员工的薪资集中在中高位，而低薪员工（人力资源部赵六，¥5,000）显著拉低了整体均值。\n",
      "\n",
      "#### 2️⃣ 薪资差距悬殊\n",
      "- 最高薪资（¥20,000）是最低薪资（¥5,000）的 **4 倍**，极差高达 **¥15,000**\n",
      "- 标准差约 ¥5,891，说明薪资分布相对分散，员工间收入差异较大\n",
      "\n",
      "#### 3️⃣ 部门间薪资不均衡\n",
      "- **技术部**平均薪资最高（约 ¥18,000），技术岗薪资领先\n",
      "- **销售部**平均薪资居中（约 ¥14,000）\n",
      "- **人力资源部**薪资最低（仅 ¥5,000），与最高部门相差 3.6 倍\n",
      "\n",
      "### 四、业务建议\n",
      "\n",
      "1. **关注薪资公平性**：人力资源部薪资远低于其他部门，建议评估是否存在岗位价值评估或薪酬定位问题\n",
      "2. **薪资结构优化**：整体薪资偏度较大，可考虑完善薪酬体系，适当提升低薪岗位的竞争力\n",
      "3. **数据限制提示**：当前仅 5 名员工，样本量较小，建议结合更多员工数据得出更可靠的结论\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")  # 抑制 matplotlib 的字体警告\n",
    "\n",
    "response = visualization_agent.invoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"分析一下员工薪资的分布情况，包括均值、中位数、最大最小值等统计量\"}]\n",
    "})\n",
    "final_msg = response[\"messages\"][-1]\n",
    "print(f\"\\n分析结果：\\n{final_msg.content}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "15deca7a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== 薪资统计概览 ===\n",
      "count        5.00000\n",
      "mean     13800.00000\n",
      "std       5890.67059\n",
      "min       5000.00000\n",
      "25%      11000.00000\n",
      "50%      16000.00000\n",
      "75%      17000.00000\n",
      "max      20000.00000\n",
      "Name: salary, dtype: float64\n",
      "\n",
      "中位数薪资：16,000 元\n",
      "薪资范围：5,000 ~ 20,000 元\n",
      "薪资标准差：5,891 元\n"
     ]
    }
   ],
   "source": [
    "# 数据探索\n",
    "print(\"=== 薪资统计概览 ===\")\n",
    "print(employees_df['salary'].describe())\n",
    "\n",
    "# 计算额外统计量\n",
    "median_salary = employees_df['salary'].median()\n",
    "print(f\"\\n中位数薪资：{median_salary:,.0f} 元\")\n",
    "print(f\"薪资范围：{employees_df['salary'].min():,.0f} ~ {employees_df['salary'].max():,.0f} 元\")\n",
    "print(f\"薪资标准差：{employees_df['salary'].std():,.0f} 元\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "055cc333",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "findfont: Failed to find font weight bold, now using 400.\n",
      "findfont: Failed to find font weight bold, now using 400.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 42888 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "分析结果：\n",
      "---\n",
      "\n",
      "## 📊 各部门平均薪资对比分析\n",
      "\n",
      "### 数据概况\n",
      "| 部门 | 平均薪资 | 员工数 |\n",
      "|------|---------|:-----:|\n",
      "| 🏆 **技术部** | **¥18,000** | 2人 |\n",
      "| 🥈 销售部 | ¥14,000 | 2人 |\n",
      "| 🥉 人力资源部 | ¥5,000 | 1人 |\n",
      "\n",
      "### 📌 业务洞察\n",
      "\n",
      "1. **技术部薪资最高**（¥18,000），分别比销售部和人力资源部高出 **28.6%** 和 **260%**，符合市场规律——技术岗通常享有较高薪酬溢价。\n",
      "\n",
      "2. **人力资源部薪资明显偏低**（¥5,000），仅为公司平均薪资（¥13,800）的 **36%**，建议关注该岗位的薪酬合理性，避免因薪资过低导致人员流失。\n",
      "\n",
      "3. **部门间薪资差距较大**，最高与最低相差 **3.6倍**，建议公司评估薪酬体系的公平性与激励性，确保既有竞争力又能留住关键人才。\n",
      "\n",
      "> 💡 如需进一步分析，比如各岗位的薪资与入职时间的关系、或按部门拆分更细的薪资分布，随时告诉我！\n"
     ]
    }
   ],
   "source": [
    "response = await visualization_agent.ainvoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"画一个柱状图，展示各部门的平均薪资对比\"}]\n",
    "})\n",
    "final_msg = response[\"messages\"][-1]\n",
    "print(f\"\\n分析结果：\\n{final_msg.content}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "42693461",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31624 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 定制化\n",
    "# 设置中文字体\n",
    "plt.rcParams['font.sans-serif'] = ['SimHei', 'PingFang SC', 'DejaVu Sans']\n",
    "plt.rcParams['axes.unicode_minus'] = False\n",
    "\n",
    "# 按部门计算平均薪资\n",
    "dept_salary = employees_df.groupby('department')['salary'].mean().sort_values(ascending=False)\n",
    "\n",
    "# 画柱状图\n",
    "plt.figure(figsize=(10, 6))\n",
    "bars = plt.bar(dept_salary.index, dept_salary.values, color=['#4C78A8', '#F58518', '#E45756'])\n",
    "plt.title('Average Salary by Department', fontsize=16)\n",
    "plt.xlabel('Department', fontsize=12)\n",
    "plt.ylabel('Average Salary (CNY)', fontsize=12)\n",
    "\n",
    "# 在柱子上方标注具体数值\n",
    "for bar in bars:\n",
    "    height = bar.get_height()\n",
    "    plt.text(bar.get_x() + bar.get_width()/2., height,\n",
    "             f'{height:,.0f}', ha='center', va='bottom', fontsize=11)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "bd05ad24",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "findfont: Failed to find font weight bold, now using 400.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 38032 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "分析结果：\n",
      "## 分析结果与业务洞察\n",
      "\n",
      "### 📊 各产品类别库存总量\n",
      "\n",
      "| 产品类别 | 库存总量 |\n",
      "|---------|:-------:|\n",
      "| 办公用品 | 300 件 |\n",
      "| 电子产品 | **1,900 件** |\n",
      "| **合计** | **2,200 件** |\n",
      "\n",
      "---\n",
      "\n",
      "### 🔍 业务洞察\n",
      "\n",
      "1. **电子产品占据绝对主导地位** — 库存总量高达 **1,900 件**，占总库存的 **86.4%**。这主要是因为产品表中有 3 个电子产品（笔记本电脑、机械键盘、显示器），而办公用品只有 1 个（办公椅）。\n",
      "\n",
      "2. **库存结构不平衡** — 电子产品类别涵盖了从 399 元的机械键盘到 6,999 元的笔记本电脑，品类丰富度更高；而办公用品仅有一款产品，库存深度明显不足。\n",
      "\n",
      "3. **潜在建议**：\n",
      "   - 若办公用品需求稳定，可考虑扩充该品类产品线，增加 SKU 数量，优化库存结构\n",
      "   - 电子产品的库存量较大，需关注周转率，避免资金被过度占用（尤其是高单价商品如笔记本电脑）\n",
      "\n",
      "> 📌 从水平条形图可以直观看出，电子产品的库存量是办公用品的 **6 倍以上**，两者差距悬殊。\n"
     ]
    }
   ],
   "source": [
    "response = await visualization_agent.ainvoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"计算每个产品类别的库存总量，用水平条形图展示，并标注具体数值\"}]\n",
    "})\n",
    "final_msg = response[\"messages\"][-1]\n",
    "print(f\"\\n分析结果：\\n{final_msg.content}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "37fef040",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50832 characters]",
      "text/plain": [
       "<Figure size 1000x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "============================================================\n",
      "Agent 生成的第 2 段代码：\n",
      "============================================================\n",
      "# 先探索产品表的结构\n",
      "print(\"=== 产品表基本信息 ===\")\n",
      "print(products_df.info())\n",
      "print(\"\\n=== 前5行数据 ===\")\n",
      "print(products_df.head())\n",
      "print(\"\\n=== 基本统计 ===\")\n",
      "print(products_df.describe())\n",
      "print(\"\\n=== 类别分布 ===\")\n",
      "print(products_df['category'].value_counts())\n",
      "\n",
      "--- 执行结果 ---\n",
      "执行成功：\n",
      "=== 产品表基本信息 ===\n",
      "<class 'pandas.DataFrame'>\n",
      "RangeIndex: 4 entries, 0 to 3\n",
      "Data columns (total 5 columns):\n",
      " #   Column        Non-Null Count  Dtype  \n",
      "---  ------        --------------  -----  \n",
      " 0   id            4 non-null      int64  \n",
      " 1   product_name  4 non-null      str    \n",
      " 2   category      4 non-null      str    \n",
      " 3   price         4 non-null      float64\n",
      " 4   stock         4 non-null      int64  \n",
      "dtypes: float64(1), int64(2), str(2)\n",
      "memory usage: 292.0 bytes\n",
      "None\n",
      "\n",
      "=== 前5行数据 ===\n",
      "   id\n",
      "\n",
      "============================================================\n",
      "Agent 生成的第 4 段代码：\n",
      "============================================================\n",
      "# 设置中文字体\n",
      "plt.rcParams['font.sans-serif'] = ['SimHei', 'PingFang SC', 'DejaVu Sans']\n",
      "plt.rcParams['axes.unicode_minus'] = False\n",
      "\n",
      "# 计算各产品类别的平均价格\n",
      "avg_price_by_category = products_df.groupby('category')['price'].mean().round(2)\n",
      "print(\"=== 各产品类别平均价格 ===\")\n",
      "print(avg_price_by_category)\n",
      "print(f\"\\n总平均价格: {products_df['price'].mean():.2f}\")\n",
      "\n",
      "# 画饼图\n",
      "plt.figure(figsize=(10, 8))\n",
      "colors = ['#ff9999', '#66b3ff', '#99ff99', '#ffcc99']\n",
      "explode = [0.05] * len(avg_price_by_category)  # 略微分离\n",
      "\n",
      "# 饼图\n",
      "wedges, texts, autotexts = plt.pie(\n",
      "    avg_price_by_category.values,\n",
      "    labels=avg_price_by_category.index,\n",
      "    autopct='%1.1f%%',\n",
      "    colors=colors[:len(avg_price_by_category)],\n",
      "    explode=explode,\n",
      "    startangle=90,\n",
      "    textprops={'fontsize': 14}\n",
      ")\n",
      "\n",
      "# 增强文字样式\n",
      "for text in autotexts:\n",
      "    text.set_fontsize(14)\n",
      "    text.set_fontweight('bold')\n",
      "\n",
      "# 添加图例显示具体数值\n",
      "legend_labels = [f'{cat}: ¥{val:,.2f}' for cat, val in zip(avg_price_by_category.index, avg_price_by_category.values)]\n",
      "plt.legend(legend_labels, title=\"类别 (平均价格)\", loc=\"upper right\", fontsize=12)\n",
      "\n",
      "plt.title('Product Category Average Price Distribution', fontsize=16, fontweight='bold', pad=20)\n",
      "plt.tight_layout()\n",
      "plt.show()\n",
      "\n",
      "--- 执行结果 ---\n",
      "执行成功：\n",
      "=== 各产品类别平均价格 ===\n",
      "category\n",
      "办公用品     499.0\n",
      "电子产品    2866.0\n",
      "Name: price, dtype: float64\n",
      "\n",
      "总平均价格: 2274.25\n",
      "\n",
      "\n",
      "============================================================\n",
      "最终回答：\n",
      "## 📊 分析结果与业务洞察\n",
      "\n",
      "### 各产品类别平均价格\n",
      "\n",
      "| 类别 | 平均价格 | 占比 |\n",
      "|:---:|:---:|:---:|\n",
      "| 🖥️ **电子产品** | **¥2,866.00** | **85.2%** |\n",
      "| 🪑 **办公用品** | **¥499.00** | **14.8%** |\n",
      "| **总体平均** | **¥2,274.25** | — |\n",
      "\n",
      "### 🔍 洞察分析\n",
      "\n",
      "1. **价格差异巨大**：电子产品的平均价格（¥2,866）是办公用品（¥499）的 **5.7 倍**，反映出电子产品本身单价高、毛利空间大的特点。\n",
      "\n",
      "2. **品类结构不均衡**：当前产品线中 **电子产品占 75%**（3/4 个产品），而办公用品只有 1 个产品。建议可适当丰富办公用品品类，均衡产品结构。\n",
      "\n",
      "3. **业务建议**：\n",
      "   - 电子产品是核心收入贡献品类，应重点维护并关注高单价产品（如笔记本电脑 ¥6,999）的库存和销售情况\n",
      "   - 办公用品目前只有 1 款产品（办公椅），可考虑扩充品类增加收入来源\n",
      "\n",
      "饼图直观展示了两个类别平均价格的比例关系，**电子产品占据了绝对主导地位**（85.2%）。\n"
     ]
    }
   ],
   "source": [
    "# 查看 Agent 的完整推理过程，重点是它生成了什么代码\n",
    "response = await visualization_agent.ainvoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"分析各产品类别的平均价格，画一个饼图\"}]\n",
    "})\n",
    "\n",
    "for i, msg in enumerate(response[\"messages\"]):\n",
    "    msg_type = msg.__class__.__name__\n",
    "\n",
    "    # AIMessage 中的 tool_calls 包含 Agent 生成的代码\n",
    "    if msg_type == \"AIMessage\" and hasattr(msg, \"tool_calls\") and msg.tool_calls:\n",
    "        print(f\"\\n{'='*60}\")\n",
    "        for tc in msg.tool_calls:\n",
    "            if tc[\"name\"] == \"execute_python_code\":\n",
    "                print(f\"Agent 生成的第 {i+1} 段代码：\")\n",
    "                print(f\"{'='*60}\")\n",
    "                print(tc[\"args\"].get(\"code\", \"\"))\n",
    "\n",
    "    # ToolMessage 是代码执行的结果\n",
    "    elif msg_type == \"ToolMessage\":\n",
    "        print(f\"\\n--- 执行结果 ---\")\n",
    "        print(msg.content[:500])\n",
    "\n",
    "# 最终回答\n",
    "final_msg = response[\"messages\"][-1]\n",
    "print(f\"\\n{'='*60}\")\n",
    "print(f\"最终回答：\\n{final_msg.content}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "db5050f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "用户意图：both\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "findfont: Failed to find font weight bold, now using 400.\n",
      "findfont: Failed to find font weight bold, now using 400.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 58252 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "最终回答：\n",
      "## 📊 分析结果与业务洞察\n",
      "\n",
      "### 数据概览\n",
      "\n",
      "| 指标 | 数值 |\n",
      "|------|:----:|\n",
      "| 员工总数 | **5 人** |\n",
      "| 薪资 > 18,000 的员工 | **1 人（20%）** |\n",
      "| 薪资 ≤ 18,000 的员工 | **4 人（80%）** |\n",
      "| 公司平均薪资 | **¥13,800** |\n",
      "| 薪资范围 | **¥5,000 ~ ¥20,000** |\n",
      "\n",
      "### 🎯 高薪员工详情\n",
      "\n",
      "| 姓名 | 部门 | 薪资 |\n",
      "|:----:|:----:|:----:|\n",
      "| **张三** | **技术部** | **¥20,000** |\n",
      "\n",
      "---\n",
      "\n",
      "### 💡 业务洞察\n",
      "\n",
      "1. **高薪员工稀缺**：薪资超过 18,000 元的员工仅 **1 人**，占全员的 **20%**，说明公司当前薪资结构偏集中，高薪段人才非常稀少。\n",
      "\n",
      "2. **技术部贡献高薪**：唯一的高薪员工 **张三** 来自 **技术部**，薪资为 **¥20,000**，是公司最高薪资，同时也是唯一超过 18,000 元门槛的员工。\n",
      "\n",
      "3. **薪资分层明显**：其余 4 名员工薪资分布在 5,000~17,000 元区间，与张三的 20,000 元有一定差距，薪资梯队断层显著。\n",
      "\n",
      "4. **潜在建议**：\n",
      "   - 若公司希望吸引更多高薪人才，可考虑调整薪酬体系，拓宽高薪段覆盖面\n",
      "   - 技术部作为高薪核心部门，可重点加大技术人才的招聘与保留投入\n",
      "   - 当前仅 20% 的员工处于高薪段，薪酬竞争力可能面临挑战\n",
      "\n",
      "> 📌 **饼状图说明**：由于实际符合条件（薪资 > 18,000）的员工仅 1 人，饼图按\"高薪段 vs 普通段\"两个区间分组展示，清晰反映了薪资分布的两极分化态势。\n"
     ]
    }
   ],
   "source": [
    "# 编排器：根据用户意图路由到合适的 Agent\n",
    "async def run_data_analysis(user_query: str) -> str:\n",
    "    \"\"\"\n",
    "    统一入口：自动判断用户意图，路由到对应的 Agent。\n",
    "    \"\"\"\n",
    "    # 用 LLM 做意图分类（也可以用规则匹配，看场景复杂度）\n",
    "    classification_prompt = f\"\"\"判断以下用户问题属于哪种类型：\n",
    "    - \"query\"：需要查询数据库获取数据\n",
    "    - \"visualize\"：需要画图或做统计分析\n",
    "    - \"both\"：需要先查数据，再画图分析\n",
    "\n",
    "    用户问题：{user_query}\n",
    "\n",
    "    只回复一个词：query / visualize / both\"\"\"\n",
    "\n",
    "    response = await llm.ainvoke(classification_prompt)\n",
    "    intent = response.content.strip().lower()\n",
    "    print(f\"用户意图：{intent}\")\n",
    "\n",
    "    # 根据 intent 路由到对应的 Agent\n",
    "    if intent == \"query\":\n",
    "        # create_agent 返回的消息格式：取最后一条消息\n",
    "        result = await sql_agent.ainvoke({\"messages\": [{\"role\": \"user\", \"content\": user_query}]})\n",
    "        return result[\"messages\"][-1].content\n",
    "    elif intent == \"visualize\":\n",
    "        result = await visualization_agent.ainvoke({\"messages\": [{\"role\": \"user\", \"content\": user_query}]})\n",
    "        return result[\"messages\"][-1].content\n",
    "    else:  # both\n",
    "        # 先查数据\n",
    "        data_result = await sql_agent.ainvoke({\"messages\": [{\"role\": \"user\", \"content\": user_query}]})\n",
    "        data_content = data_result[\"messages\"][-1].content\n",
    "        # 把查询结果传给可视化 Agent\n",
    "        viz_input = f\"基于以下数据进行可视化分析：\\n{data_content}\\n\\n原始问题：{user_query}\"\n",
    "        viz_result = await visualization_agent.ainvoke({\"messages\": [{\"role\": \"user\", \"content\": viz_input}]})\n",
    "        return viz_result[\"messages\"][-1].content\n",
    "\n",
    "resp = await run_data_analysis(\"用饼状图展示薪资大于 18000 的员工\")\n",
    "print(f\"最终回答：\\n{resp}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f666e0b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "【纯SQL】Agent 工具调用链：\n",
      "  [AIMessage] → sql_db_list_tables({})\n",
      "  [AIMessage] → sql_db_schema({'table_names': 'employees'})\n",
      "  [AIMessage] → sql_db_query_checker({'query': 'SELECT department, COUNT(*) AS 人数 FROM employees GROUP BY department ORDER BY 人数 DESC'})\n",
      "  [AIMessage] → sql_db_query({'query': 'SELECT department, COUNT(*) AS 人数 FROM employees GROUP BY department ORDER BY 人数 DESC'})\n",
      "回答： 以下是公司各部门人数统计结果：\n",
      "\n",
      "| 部门 | 人数 |\n",
      "|------|:----:|\n",
      "| **技术部** | 2 人 |\n",
      "| **销售部** | 2 人 |\n",
      "| **人力资源** | 1 人 |\n",
      "\n",
      "**公司总人数：5 人**\n",
      "\n",
      "📌 **小结：**\n",
      "- 技术部和销售部人数相同，均为 **2 人**，各占公司总人数的 **40%**。\n",
      "- 人力资源部目前只有 **1 人**，占比 **20%**。\n",
      "\n",
      "【纯可视化】Agent 工具调用链：\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 43880 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  [AIMessage] → sql_db_list_tables({})\n",
      "  [AIMessage] → sql_db_schema({'table_names': 'employees'})\n",
      "  [AIMessage] → execute_python_code({'code': '# 设置字体\\nplt.rcParams[\\'font.sans-serif\\'] = [\\'SimHei\\', \\'PingFang SC\\', \\'DejaVu Sans\\']\\nplt.rcParams[\\'axes.unicode_minus\\'] = False\\n\\n# 查看员工数据\\nprint(employees_df.head())\\nprint(\"\\\\n--- 各部门平均薪资 ---\")\\navg_salary = employees_df.groupby(\\'department\\')[\\'salary\\'].mean().sort_values(ascending=False)\\nprint(avg_salary)\\n\\n# 画柱状图\\nplt.figure(figsize=(10, 6))\\ncolors = [\\'#3498db\\', \\'#e74c3c\\', \\'#2ecc71\\', \\'#f39c12\\', \\'#9b59b6\\']\\nbars = plt.bar(avg_salary.index, avg_salary.values, color=colors[:len(avg_salary)], edgecolor=\\'white\\', linewidth=1.2)\\n\\n# 在柱子上方标注数值\\nfor bar, val in zip(bars, avg_salary.values):\\n    plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 200, \\n             f\\'¥{val:,.0f}\\', ha=\\'center\\', va=\\'bottom\\', fontsize=12, fontweight=\\'bold\\')\\n\\nplt.title(\\'各部门平均薪资对比\\', fontsize=16, fontweight=\\'bold\\', pad=15)\\nplt.xlabel(\\'部门\\', fontsize=13)\\nplt.ylabel(\\'平均薪资（元）\\', fontsize=13)\\nplt.xticks(fontsize=12)\\nplt.yticks(fontsize=11)\\nplt.grid(axis=\\'y\\', alpha=0.3, linestyle=\\'--\\')\\nplt.tight_layout()\\nplt.show()'})\n",
      "回答： ## 各部门平均薪资对比\n",
      "\n",
      "从数据来看，各部门平均薪资差异明显：\n",
      "\n",
      "| 部门 | 平均薪资 |\n",
      "|------|---------|\n",
      "| 🔧 **技术部** | **¥18,000** |\n",
      "| 💼 **销售部** | **¥14,000** |\n",
      "| 👥 **人力资源** | **¥5,000** |\n",
      "\n",
      "### 结论\n",
      "- **技术部**薪资最高（¥18,000），是人力资源部的 **3.6 倍**；\n",
      "\n",
      "【混合场景】Agent 工具调用链：\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 38708 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  [AIMessage] → sql_db_list_tables({})\n",
      "  [AIMessage] → sql_db_schema({'table_names': 'employees, orders, products'})\n",
      "  [AIMessage] → sql_db_query({'query': \"SELECT e.id, e.name, e.department, \\n       SUM(o.quantity * p.price) AS total_amount\\nFROM employees e\\nJOIN orders o ON e.id = o.employee_id\\nJOIN products p ON o.product_id = p.id\\nWHERE e.department = '销售部'\\nGROUP BY e.id, e.name, e.department\\nORDER BY total_amount DESC\"})\n",
      "  [AIMessage] → sql_db_query({'query': \"SELECT id, name, department FROM employees WHERE department = '销售部'\"})\n",
      "  [AIMessage] → execute_python_code({'code': \"plt.rcParams['font.sans-serif'] = ['SimHei', 'PingFang SC', 'DejaVu Sans']\\nplt.rcParams['axes.unicode_minus'] = False\\n\\n# 数据\\ndata = {\\n    'name': ['李四', '钱七'],\\n    'total_amount': [11985.00, 2994.00]\\n}\\n\\ndf = pd.DataFrame(data)\\n\\n# 画柱状图\\nplt.figure(figsize=(10, 6))\\ncolors = ['#4C72B0', '#DD8452']\\nbars = plt.bar(df['name'], df['total_amount'], color=colors, width=0.5, edgecolor='gray')\\n\\n# 在柱子上标数值\\nfor bar in bars:\\n    height = bar.get_height()\\n    plt.text(bar.get_x() + bar.get_width()/2., height + 100,\\n             f'¥{height:,.0f}', ha='center', va='bottom', fontsize=13, fontweight='bold')\\n\\nplt.title('销售部员工订单总金额对比', fontsize=16, fontweight='bold')\\nplt.xlabel('员工姓名', fontsize=13)\\nplt.ylabel('订单总金额 (元)', fontsize=13)\\nplt.ylim(0, 15000)\\nplt.grid(axis='y', linestyle='--', alpha=0.6)\\n\\nplt.tight_layout()\\nplt.show()\"})\n",
      "回答： ## 📊 销售部员工订单总金额对比\n",
      "\n",
      "| 员工 | 订单总金额 |\n",
      "|------|:---------:|\n",
      "| **李四** | **¥11,985** |\n",
      "| **钱七** | **¥2,994** |\n",
      "\n",
      "### 📌 分析结论\n",
      "\n",
      "1. **李四** 订单总金额为 **¥11,985**，占销售部总额的 **80%**，是销售主力。\n",
      "2. **钱七** 订单总金额为 **¥2,994**，占销售\n"
     ]
    }
   ],
   "source": [
    "# 方案B：统一Agent —— SQL工具 + Python执行工具 挂一起\n",
    "unified_tools = tools + [execute_python_code]\n",
    "\n",
    "UNIFIED_AGENT_PROMPT = \"\"\"\\\n",
    "你是一名数据分析师，会写SQL也会用Python做分析和画图。\n",
    "\n",
    "工具：\n",
    "- sql_db_list_tables / sql_db_schema / sql_db_query_checker / sql_db_query：查数据库\n",
    "- execute_python_code：执行Python代码，可用变量 employees_df, products_df, orders_df, pd, plt, sns, np\n",
    "\n",
    "规则：\n",
    "- 查库：先list_tables → 看schema → 检查SQL → 执行\n",
    "- 画图：用execute_python_code，绑图前设 plt.rcParams['font.sans-serif'] = ['SimHei', 'PingFang SC', 'DejaVu Sans'] 和 plt.rcParams['axes.unicode_minus'] = False，图尺寸统一 plt.figure(figsize=(10,6))，加标题和轴标签\n",
    "- 回答用中文，简洁有数据有结论\n",
    "\"\"\"\n",
    "\n",
    "unified_agent = create_agent(model=llm, tools=unified_tools, system_prompt=UNIFIED_AGENT_PROMPT)\n",
    "\n",
    "def show_tool_chain(resp):\n",
    "    \"\"\"打印 Agent 调用过的工具链\"\"\"\n",
    "    for msg in resp[\"messages\"]:\n",
    "        t = msg.__class__.__name__\n",
    "        if hasattr(msg, \"tool_calls\") and msg.tool_calls:\n",
    "            for tc in msg.tool_calls:\n",
    "                print(f\"  [{t}] → {tc['name']}({tc.get('args', {})})\")\n",
    "\n",
    "# ---- 测试：纯SQL ----\n",
    "print(\"【纯SQL】Agent 工具调用链：\")\n",
    "resp = await unified_agent.ainvoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"公司每个部门各有多少人？\"}]\n",
    "})\n",
    "show_tool_chain(resp)\n",
    "print(\"回答：\", resp[\"messages\"][-1].content[:200])\n",
    "\n",
    "# ---- 测试：纯可视化 ----\n",
    "print(\"\\n【纯可视化】Agent 工具调用链：\")\n",
    "resp = await unified_agent.ainvoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"画柱状图，展示各部门平均薪资对比\"}]\n",
    "})\n",
    "show_tool_chain(resp)\n",
    "print(\"回答：\", resp[\"messages\"][-1].content[:200])\n",
    "\n",
    "# ---- 测试：混合（先查库再画图） ----\n",
    "print(\"\\n【混合场景】Agent 工具调用链：\")\n",
    "resp = await unified_agent.ainvoke({\n",
    "    \"messages\": [{\"role\": \"user\", \"content\": \"查询销售部每位员工的订单总金额，用柱状图对比\"}]\n",
    "})\n",
    "show_tool_chain(resp)\n",
    "print(\"回答：\", resp[\"messages\"][-1].content[:200])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "01_langchain (3.11.14.final.0)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
