agentlearning_repo/01_dataanalyst/agents/viz_agent.py

"""
数据可视化 Agent：用 Python 代码对 DataFrame 做统计分析和绑图。
"""

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

# 确保父目录可导入（支持直接运行 python agents/viz_agent.py）
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from langchain.tools import tool
from langchain.agents import create_agent

from config import llm, engine

warnings.filterwarnings("ignore")

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

# ============================================================
# 2. 延迟加载 DataFrame（避免 import 时就连库）
# ============================================================
_dataframes = None


def load_dataframes():
    """从数据库加载 DataFrame。首次调用后缓存，后续直接返回缓存。"""
    global _dataframes
    if _dataframes is None:
        _dataframes = {
            "employees_df": pd.read_sql("SELECT * FROM employees", engine),
            "products_df":  pd.read_sql("SELECT * FROM products", engine),
            "orders_df":    pd.read_sql("SELECT * FROM orders", engine),
        }
    return _dataframes


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

    可用变量：
      - employees_df: 员工表 DataFrame（字段：id, name, department, salary, hire_date）
      - products_df:  产品表 DataFrame（字段：id, product_name, category, price, stock）
      - orders_df:    订单表 DataFrame（字段：id, employee_id, product_id, quantity, order_date）
      - pd:   pandas 库
      - plt:  matplotlib.pyplot
      - sns:  seaborn
      - np:   numpy

    使用示例：
      result = employees_df.groupby('department')['salary'].mean()
      print(result)
    """
    dfs = load_dataframes()
    sandbox = {
        **dfs,
        "pd":  pd,
        "plt": plt,
        "sns": sns,
        "np":  np,
    }
    exec_locals = {}
    output_buffer = StringIO()

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


# ============================================================
# 4. Agent System Prompt
# ============================================================
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"
"""

# ============================================================
# 5. 创建 Agent
# ============================================================
visualization_agent = create_agent(
    model=llm,
    tools=[execute_python_code],
    system_prompt=VISUALIZATION_PROMPT,
)


if __name__ == "__main__":
    dfs = load_dataframes()
    print("✅ 数据加载完成")
    for name, df in dfs.items():
        print(f"   {name}：{len(df)} 行 × {len(df.columns)} 列")
    print("\n✅ 数据可视化 Agent 创建完成")