assignment/01_data_analysis.py

import os
import traceback
from io import StringIO
from contextlib import redirect_stdout

import mysql.connector
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from dotenv import load_dotenv
from langchain.tools import tool
from langchain.agents import create_agent
from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from sqlalchemy import create_engine


class DataAnalysisSystem:
    """
    数据分析系统：集成 NL2SQL、数据可视化、意图路由的统一分析平台。
    支持硅基流动/DeepSeek 等大模型 API。
    """

    # ============================================================
    # 类级常量：Prompt 模板
    # ============================================================
    SQL_AGENT_PROMPT = """你是一名专业的 SQL 数据分析师。

        ## 工作流程
        1. 先用 sql_db_list_tables 查看数据库中有哪些表
        2. 用 sql_db_schema 获取相关表的字段结构和类型
        3. 生成 SQL 之前，用 sql_db_query_checker 检查语法
        4. 确认无误后，用 sql_db_query 执行查询
        5. 用中文总结查询结果，给出简洁的业务洞察
        
        ## 约束
        - 只使用数据库中实际存在的表和字段，不要凭空编造
        - 单次查询结果限制在 50 条以内
        - 如果查询出错，分析错误原因后重新生成 SQL
        - 回答要简洁专业，不要啰嗦
        """

    VISUALIZATION_PROMPT = """你是一名资深数据分析师，精通 Python、Pandas 和 Matplotlib 数据可视化。

        ## 可用数据
        1. employees_df — 员工表（字段：id, name, department, salary, hire_date）
        2. products_df  — 产品表（字段：id, product_name, category, price, stock）
        3. orders_df    — 订单表（字段：id, employee_id, product_id, quantity, order_date）
        
        ## 工作流程
        1. 理解用户的分析需求
        2. 用 execute_python_code 工具编写并执行 Python 代码
        3. 先做数据探索（head、describe、info），再做深入分析
        4. 用中文解释分析结果，给出业务洞察

        ## 代码规范
        - 绑图前设置中文字体：plt.rcParams['font.sans-serif'] = ['SimHei', 'PingFang SC', 'DejaVu Sans']
        - 设置 plt.rcParams['axes.unicode_minus'] = False
        - 图表尺寸统一用 plt.figure(figsize=(10, 6))
        - 必须添加标题、坐标轴标签，让图表自解释
        - 用 print() 输出关键统计量，不要只画图不说话
        - 图表标题用英文（避免渲染问题），但用中文向用户解释结果
        
        ## 注意事项
        - 每次只执行一段完整的代码，不要拆成多段
        - 先探索数据结构，再做分析——不要上来就画图
        - 结果要有业务洞察，不只是"最大值是 XXX"
        """

    # ============================================================
    # 构造函数
    # ============================================================
    def __init__(self, env_file: str = ".env"):
        """
        初始化数据分析系统。

        Args:
            env_file: 环境变量文件路径，默认当前目录的 .env
        """
        # 加载环境变量
        load_dotenv(dotenv_path=env_file)

        # 初始化 LLM
        self.api_key = os.getenv("API_KEY")
        self.model_name = os.getenv("MODEL_NAME")
        self.base_url = os.getenv("BASE_URL")

        self.llm = ChatOpenAI(
            model=self.model_name,
            api_key=self.api_key,
            base_url=self.base_url,
            temperature=0,
        )
        print("✅ 大语言模型初始化完成")

        # 构建数据库 URI
        self.db_uri = (
            f"mysql+pymysql://{os.getenv('DB_USER')}:{os.getenv('DB_PASSWORD')}"
            f"@{os.getenv('DB_HOST')}:{os.getenv('DB_PORT')}/{os.getenv('DB_NAME')}"
        )
        print("✅ 数据库 URI 构建完成")

        # 数据缓存
        self.employees_df = None
        self.products_df = None
        self.orders_df = None

        # Agent 缓存（延迟初始化）
        self._sql_agent = None
        self._visual_agent = None
        self._sandbox_globals = None

        self.together_agent = None


    # ============================================================
    # 公共方法：测试模型连接
    # ============================================================
    def test_llm(self, prompt: str = "什么是量子纠缠？") -> str:
        """测试 LLM 连接是否正常。"""
        response = self.llm.invoke(prompt)
        return response.content

    # ============================================================
    # 公共方法：初始化数据库（建表 + 插数据）
    # ============================================================
    def init_database(self) -> None:
        """
        创建演示数据库：employees、products、orders 三张表，
        并插入示例数据。
        """
        conn = mysql.connector.connect(
            host=os.getenv("DB_HOST"),
            port=int(os.getenv("DB_PORT")),
            user=os.getenv("DB_USER"),
            password=os.getenv("DB_PASSWORD"),
            database=os.getenv("DB_NAME"),
        )
        cursor = conn.cursor()

        # 建表
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS employees (
            id          INT PRIMARY KEY AUTO_INCREMENT,
            name        VARCHAR(50)  NOT NULL,
            department  VARCHAR(50)  NOT NULL,
            salary      DECIMAL(10,2) NOT NULL,
            hire_date   DATE
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
        """)

        cursor.execute("""
        CREATE TABLE IF NOT EXISTS products (
            id           INT PRIMARY KEY AUTO_INCREMENT,
            product_name VARCHAR(100) NOT NULL,
            category     VARCHAR(50)  NOT NULL,
            price        DECIMAL(10,2) NOT NULL,
            stock        INT DEFAULT 0
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
        """)

        cursor.execute("""
        CREATE TABLE IF NOT EXISTS orders (
            id           INT PRIMARY KEY AUTO_INCREMENT,
            employee_id  INT NOT NULL,
            product_id   INT NOT NULL,
            quantity     INT NOT NULL,
            order_date   DATE NOT NULL,
            FOREIGN KEY (employee_id) REFERENCES employees(id),
            FOREIGN KEY (product_id)  REFERENCES products(id)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
        """)

        # 插入数据
        employees_data = [
            (1, "张三", "技术部", 20000.00, "2023-01-15"),
            (2, "李四", "销售部", 11000.00, "2023-02-20"),
            (3, "王五", "技术部", 16000.00, "2022-11-10"),
            (4, "赵六", "人力资源", 5000.00, "2023-03-01"),
            (5, "钱七", "销售部", 17000.00, "2022-12-05"),
        ]
        products_data = [
            (1, "笔记本电脑", "电子产品", 6999.00, 500),
            (2, "机械键盘", "电子产品", 399.00, 1000),
            (3, "办公椅", "办公用品", 499.00, 300),
            (4, "显示器", "电子产品", 1200.00, 400),
        ]
        orders_data = [
            (1, 1, 1, 2, "2024-01-15"),
            (2, 2, 2, 15, "2024-01-16"),
            (3, 3, 1, 10, "2024-01-17"),
            (4, 5, 3, 6, "2024-01-18"),
            (5, 2, 4, 5, "2024-01-19"),
        ]

        cursor.executemany(
            "INSERT IGNORE INTO employees VALUES (%s,%s,%s,%s,%s)", employees_data
        )
        cursor.executemany(
            "INSERT IGNORE INTO products  VALUES (%s,%s,%s,%s,%s)", products_data
        )
        cursor.executemany(
            "INSERT IGNORE INTO orders     VALUES (%s,%s,%s,%s,%s)", orders_data
        )

        conn.commit()
        conn.close()

        print("✅ 数据库初始化完成")
        print("   - employees 表：5 条员工记录")
        print("   - products  表：4 条产品记录")
        print("   - orders    表：5 条订单记录")

    # ============================================================
    # 内部方法：获取 SQL Agent（延迟初始化）
    # ============================================================
    def _get_sql_agent(self):
        """延迟初始化 SQL Agent。"""
        if self._sql_agent is not None:
            return self._sql_agent

        db = SQLDatabase.from_uri(self.db_uri)
        print(f"数据库连接成功，可用表：{db.get_usable_table_names()}")

        toolkit = SQLDatabaseToolkit(db=db, llm=self.llm)
        tools = toolkit.get_tools()

        print(f"SQL 工具包已加载（{len(tools)} 个工具）：")
        for t in tools:
            print(f"   - {t.name}")

        self._sql_agent = create_agent(
            model=self.llm,
            tools=tools,
            system_prompt=self.SQL_AGENT_PROMPT,
        )
        print("NL2SQL Agent 创建完成")
        return self._sql_agent

    # ============================================================
    # 内部方法：加载数据到 DataFrame（用于可视化）
    # ============================================================
    def _load_dataframes(self):
        """从数据库加载数据到 Pandas DataFrame。"""
        if self.employees_df is not None:
            return  # 已加载，跳过

        # 配置 matplotlib 中文显示
        plt.rcParams["font.sans-serif"] = ["SimHei", "PingFang SC", "DejaVu Sans"]
        plt.rcParams["axes.unicode_minus"] = False

        engine = create_engine(self.db_uri)

        self.employees_df = pd.read_sql("SELECT * FROM employees", engine)
        self.products_df = pd.read_sql("SELECT * FROM products", engine)
        self.orders_df = pd.read_sql("SELECT * FROM orders", engine)

        print("✅ 数据加载完成")
        print(f"   员工表：{len(self.employees_df)} 行")
        print(f"   产品表：{len(self.products_df)} 行")
        print(f"   订单表：{len(self.orders_df)} 行")

    # ============================================================
    # 内部方法：获取可视化 Agent（延迟初始化）
    # ============================================================
    def _get_visual_agent(self):
        """延迟初始化可视化 Agent。"""
        if self._visual_agent is not None:
            return self._visual_agent

        self._load_dataframes()

        # 构建沙箱全局变量
        self._sandbox_globals = {
            "employees_df": self.employees_df,
            "products_df": self.products_df,
            "orders_df": self.orders_df,
            "pd": pd,
            "plt": plt,
            "sns": sns,
            "np": np,
        }

        # 创建工具
        @tool
        def execute_python_code(code: str) -> str:
            """
            执行 Python 代码进行数据分析和可视化。

            可用变量：
              - employees_df, products_df, orders_df
              - pd, plt, sns, np
            """
            exec_globals = dict(self._sandbox_globals)
            exec_locals = {}
            output_buffer = StringIO()

            try:
                with redirect_stdout(output_buffer):
                    exec(code, exec_globals, exec_locals)
                result = output_buffer.getvalue()
                if not result.strip():
                    result = "✅ 代码执行成功（无文本输出，可能已生成图表）"
                return f"执行成功：\n{result}"
            except Exception as e:
                error_detail = traceback.format_exc()
                return f"❌ 执行出错：{e}\n\n{error_detail}"

        viz_tools = [execute_python_code]

        self._visual_agent = create_agent(
            model=self.llm,
            tools=viz_tools,
            system_prompt=self.VISUALIZATION_PROMPT,
        )
        print("✅ 数据可视化 Agent 创建完成")
        return self._visual_agent

    # ============================================================
    # 公共方法：运行 SQL 查询
    # ============================================================
    def query(self, user_query: str) -> str:
        """
        使用 SQL Agent 执行自然语言到 SQL 的查询。

        Args:
            user_query: 用户的自然语言问题

        Returns:
            Agent 的分析结果文本
        """
        agent = self._get_sql_agent()
        result = agent.invoke({"messages": [{"role": "user", "content": user_query}]})
        return result["messages"][-1].content

    # ============================================================
    # 公共方法：运行可视化分析
    # ============================================================
    def visualize(self, user_query: str) -> str:
        """
        使用可视化 Agent 进行数据分析和图表生成。

        Args:
            user_query: 用户的分析需求描述

        Returns:
            Agent 的分析结果文本
        """
        agent = self._get_visual_agent()
        result = agent.invoke({"messages": [{"role": "user", "content": user_query}]})
        return result["messages"][-1].content

    # ============================================================
    # 公共方法：统一入口（自动意图路由）
    # ============================================================
    def analyze(self, user_query: str) -> str:
        """
        统一分析入口：自动判断用户意图，路由到对应 Agent。

        意图分类：
          - query：纯数据库查询
          - visualize：纯数据可视化/统计分析
          - both：先查询再可视化

        Args:
            user_query: 用户的自然语言问题

        Returns:
            分析结果文本
        """
        classification_prompt = f"""判断以下用户问题属于哪种类型：
            - "query"：需要查询数据库获取数据
            - "visualize"：需要画图或做统计分析
            - "both"：需要先查数据，再画图分析

            用户问题：{user_query}

            只回复一个词：query / visualize / both"""

        intent = self.llm.invoke(classification_prompt).content.strip().lower()

        if intent == "query":
            return self.query(user_query)
        elif intent == "visualize":
            return self.visualize(user_query)
        else:  # both
            data_result = self.query(user_query)
            viz_input = f"基于以下数据进行可视化分析：\n{data_result}\n\n原始问题：{user_query}"
            return self.visualize(viz_input)

    # ============================================================
    # 内部方法：获取或创建统一 Agent
    # ============================================================
    def get_together_agent(self):
        TOGETHER_SYSTEM_PROMPT = """你是一名全能数据分析师，精通 SQL 查询和 Python 数据分析可视化。
               ## 可用工具
               1. **SQL 工具组**（数据库操作）：
                  - `sql_db_list_tables` — 查看数据库中有哪些表
                  - `sql_db_schema` — 获取指定表的字段结构
                  - `sql_db_query_checker` — 检查 SQL 语法是否正确
                  - `sql_db_query` — 执行 SQL 查询并返回结果

               2. **Python 代码执行工具**（数据分析与可视化）：
                  - `execute_python_code` — 在沙箱中执行 Python 代码

               ## 可用数据（Python 沙箱中已预加载）
               - `employees_df` — 员工表（id, name, department, salary, hire_date）
               - `products_df` — 产品表（id, product_name, category, price, stock）
               - `orders_df` — 订单表（id, employee_id, product_id, quantity, order_date）
               - `pd` (pandas), `plt` (matplotlib), `sns` (seaborn), `np` (numpy)

               ## 工作流程
               1. **理解需求**：分析用户问题的核心诉求
               2. **获取元数据**：如需查数据库，先用 `sql_db_list_tables` 和 `sql_db_schema` 了解表结构
               3. **选择工具**：
                  - 纯数据查询 → 用 SQL 工具
                  - 统计分析、画图 → 用 `execute_python_code`
                  - 复杂分析 → 先用 SQL 查数据，再用 Python 做深度分析
               4. **执行与验证**：SQL 先用 `sql_db_query_checker` 检查语法；Python 代码确保正确
               5. **总结输出**：用中文给出简洁的业务洞察

               ## 约束
               - 只使用数据库中实际存在的表和字段
               - SQL 单次查询结果限制在 50 条以内
               - Python 代码绑图前设置中文字体：`plt.rcParams['font.sans-serif'] = ['SimHei', 'PingFang SC', 'DejaVu Sans']`
               - 设置 `plt.rcParams['axes.unicode_minus'] = False`
               - 图表尺寸用 `plt.figure(figsize=(10, 6))`
               - 必须添加标题、坐标轴标签
               - 用 `print()` 输出关键统计量
               - 回答简洁专业，不要啰嗦
               """
        """延迟初始化统一 Agent，合并 SQL 工具 + Python 执行工具。"""
        if self.together_agent is not None:
            return self.together_agent

        # ========== 1. SQL 工具组 ==========
        db = SQLDatabase.from_uri(self.db_uri)
        print(f"数据库连接成功，可用表：{db.get_usable_table_names()}")

        sql_toolkit = SQLDatabaseToolkit(db=db, llm=self.llm)
        sql_tools = sql_toolkit.get_tools()

        print(f"SQL 工具加载完成（{len(sql_tools)} 个）：")
        for t in sql_tools:
            print(f"   - {t.name}")

        # ========== 2. Python 代码执行工具 ==========
        self._load_dataframes()

        self._sandbox_globals = {
            "employees_df": self.employees_df,
            "products_df": self.products_df,
            "orders_df": self.orders_df,
            "pd": pd,
            "plt": plt,
            "sns": sns,
            "np": np,
        }

        @tool
        def execute_python_code(code: str) -> str:
            """
            执行 Python 代码进行数据分析和可视化。

            可用变量：employees_df, products_df, orders_df, pd, plt, sns, np
            """
            exec_globals = dict(self._sandbox_globals)
            exec_locals = {}
            output_buffer = StringIO()

            try:
                with redirect_stdout(output_buffer):
                    exec(code, exec_globals, exec_locals)
                result = output_buffer.getvalue()
                if not result.strip():
                    result = "✅ 代码执行成功（无文本输出，可能已生成图表）"
                return f"执行成功：\n{result}"
            except Exception as e:
                error_detail = traceback.format_exc()
                return f"❌ 执行出错：{e}\n\n{error_detail}"

        # ========== 3. 合并所有工具 ==========
        all_tools = sql_tools + [execute_python_code]

        print(f"\n✅ 工具合并完成，共 {len(all_tools)} 个工具：")
        for t in all_tools:
            print(f"   - {t.name}")

        # ========== 4. 创建统一 Agent ==========
        self.together_agent = create_agent(
            model=self.llm,
            tools=all_tools,
            system_prompt=TOGETHER_SYSTEM_PROMPT,
        )
        print("\n🤖 统一分析 Agent 创建完成！")
        return self.together_agent

    # ============================================================
    # 统⼀ Agent 模式(作业)
    # ============================================================
    def together_analyze(self, user_query: str) -> str:
        """
        统一分析入口：LLM 自主决定调用 SQL 工具或 Python 工具。

        Args:
            user_query: 用户的自然语言问题

        Returns:
            Agent 的分析结果
        """
        agent = self.get_together_agent()
        result = agent.invoke({"messages": [{"role": "user", "content": user_query}]})
        return result["messages"][-1].content



# ============================================================
# 使用示例
# ============================================================
if __name__ == "__main__":
    # 1. 创建系统实例
    system = DataAnalysisSystem()

    # 2. 测试 LLM 连接
    print("\n--- LLM 测试 ---")
    print(system.test_llm("你好，请用一句话介绍自己"))

    # 3. 初始化数据库（只需执行一次）
    # print("\n--- 初始化数据库 ---")
    # system.init_database()

    # 4. 运行分析
    # print("\n--- SQL 查询 ---")
    # result = system.query("技术部的平均工资是多少？")
    # print(result)
    #
    # print("\n--- 可视化分析 ---")
    # result = system.visualize("统计各部门人数并画柱状图")
    # print(result)

    # print("\n--- 自动路由 ---")
    # result = system.analyze("销售部谁的总销售额最高？")
    # print(result)

    print("\n--- 统一分析：LLM 自动选择工具 ---")
    result = system.together_analyze("统计各部门人数并画柱状图？")
    print(result)