AI 驱动数据分析引擎：用 LangChain + Ollama 构建智能 EDA 助手

pip install langchain langchain-community pandas matplotlib seaborn openpyxl ollama
ollama pull llama3:latest

ai_analyst/
├── analyst.py          # 核心引擎
├── prompts.py          # 提示词模板
└── data/sample.csv     # 待分析数据集

# prompts.py
from langchain_core.prompts import ChatPromptTemplate

CODE_GENERATOR_PROMPT = ChatPromptTemplate.from_messages([
    ("system", """你是一名资深数据分析工程师。用户提供 DataFrame 变量名为 `df`，请根据需求生成纯 Python 代码。
    严格遵循：
    1. 仅允许使用 pandas、numpy、matplotlib.pyplot、seaborn 库
    2. 禁止使用 os、subprocess、eval、exec、open 等危险函数
    3. 若需绘图，必须调用 plt.show() 或 plt.savefig()
    4. 只返回代码块，不要包含任何解释性文字"""),
    ("human", "当前数据集列名：{columns}\n前5行预览：{head}\n用户需求：{query}")
])

from langchain_core.output_parsers import StrOutputParser
from langchain_community.llms import Ollama

llm = Ollama(model="llama3", temperature=0.1, num_predict=2048)
chain = CODE_GENERATOR_PROMPT | llm | StrOutputParser()

import pandas as pd
import matplotlib.pyplot as plt
import io
import sys

class AIAnalyst:
    def __init__(self, file_path: str):
        self.df = pd.read_csv(file_path)
        self.llm = Ollama(model="llama3", temperature=0)
        
    def _safe_exec(self, code: str, return_df: bool = False):
        # 构建安全的执行环境
        safe_globals = {
            'pd': pd, 'plt': plt, 'np': __import__('numpy'),
            'seaborn': __import__('seaborn'),
            '__builtins__': {'print': print, 'range': range, 'len': len, 'int': int, 'float': float, 'str': str}
        }
        local_scope = {'df': self.df.copy()}
        try:
            exec(code, safe_globals, local_scope)
            if return_df and 'df' in local_scope:
                return local_scope['df']
            return True
        except Exception as e:
            return str(e)

    def query(self, user_question: str, return_df: bool = False):
        """对外暴露的查询接口"""
        columns = ', '.join(self.df.columns)
        head = self.df.head(3).to_string()
        # 调用第一步的 chain 生成代码
        generated_code = chain.invoke({
            'columns': columns,
            'head': head,
            'query': user_question
        })
        # 清理 markdown 标记（如有）
        clean_code = generated_code.replace('```python', '').replace('```', '').strip()
        print(f"[AI 生成代码]\n{clean_code}")
        
        result = self._safe_exec(clean_code, return_df)
        if return_df and isinstance(result, pd.DataFrame):
            return result
        return result

def auto_insights(self) -> str:
    # 第一步：生成统计摘要代码
    stats_code = chain.invoke({
        'columns': ', '.join(self.df.columns),
        'head': self.df.head(2).to_string(),
        'query': '计算所有数值列的均值、中位数、标准差、最大值、最小值，并将结果保存为 stats_df 变量'
    })
    clean_stats = stats_code.replace('```python', '').replace('```', '').strip()
    
    # 执行统计代码（本地作用域捕获 stats_df）
    local_scope = {'df': self.df, 'pd': pd, 'np': __import__('numpy')}
    exec(clean_stats, {'pd': pd, 'np': __import__('numpy'), '__builtins__': {}}, local_scope)
    stats_df = local_scope.get('stats_df', pd.DataFrame())
    
    # 第二步：将统计表传给 LLM，生成业务洞察
    insight_prompt = f"以下为数据集的统计汇总：\n{stats_df.to_string()}\n请用中文总结 3 条核心业务洞察，并指出潜在异常点。"
    insight = self.llm.invoke(insight_prompt)
    return insight

def plot(self, plot_query: str, save_path: str = 'output.png'):
    query_with_path = f"{plot_query}，将图形保存为 {save_path}，dpi=150"
    code = chain.invoke({
        'columns': ', '.join(self.df.columns),
        'head': self.df.head(2).to_string(),
        'query': query_with_path
    })
    clean_code = code.replace('```python', '').replace('```', '').strip()
    self._safe_exec(clean_code)
    print(f"图表已生成：{save_path}")
    return save_path

# main.py
from analyst import AIAnalyst

analyst = AIAnalyst('data/sample.csv')

# 测试 1：数据清洗
analyst.query("将金额列中缺失值填充为平均值，数量列小于0的替换为1")

# 测试 2：复杂分析
result_df = analyst.query("按地区分组，计算每个地区的总金额和平均数量，返回结果", return_df=True)
print(result_df)

# 测试 3：自动洞察
print(analyst.auto_insights())

# 测试 4：可视化
analyst.plot("绘制各品类订单数量的箱线图，按地区区分颜色")

df['金额'].fillna(df['金额'].mean(), inplace=True)
df['数量'] = df['数量'].apply(lambda x: 1 if x < 0 else x)

# 自愈重试示例
for attempt in range(2):
    res = self._safe_exec(clean_code)
    if isinstance(res, str) and "Error" in res:
        clean_code = self.llm.invoke(f"以下代码报错：{res}，请修正：{clean_code}")
    else:
        break