如何基于LangChain的SQL Agent实现查询结果可视化？

方法1：自定义可视化工具（推荐，更可控） #

直接创建集成SQL查询与绘图逻辑的自定义工具，避免LLM直接操作Python REPL带来的风险，同时精准处理可视化需求。

步骤1：定义自定义绘图工具 #

from langchain.tools import BaseTool
from typing import Optional
from langchain.callbacks.manager import CallbackManagerForToolRun
import matplotlib.pyplot as plt
import pandas as pd
from langchain.sql_database import SQLDatabase

class SQLPlotTool(BaseTool):
    name = "sql_plot_tool"
    description = "当用户需要绘制柱状图、折线图等可视化图表时使用，输入格式为：'SQL: [你的SQL查询语句]; PLOT_TYPE: [图表类型]'，例如'SQL: SELECT category, COUNT(*) FROM products GROUP BY category; PLOT_TYPE: bar chart'"

    db: SQLDatabase

    def _run(
        self, input_str: str, run_manager: Optional[CallbackManagerForToolRun] = None
    ) -> str:
        # 解析输入内容
        parts = input_str.split("; PLOT_TYPE: ")
        if len(parts) != 2:
            return "输入格式错误，请遵循'SQL: [查询语句]; PLOT_TYPE: [图表类型]'的格式"
        
        sql_query = parts[0].replace("SQL: ", "").strip()
        plot_type = parts[1].strip().lower()

        # 执行SQL查询并转换为DataFrame
        try:
            df = pd.read_sql(sql_query, self.db._engine)
        except Exception as e:
            return f"执行SQL查询失败：{str(e)}"

        # 根据图表类型生成可视化
        plt.figure(figsize=(10, 6))
        if plot_type == "bar chart":
            df.plot(kind="bar", rot=45)
        elif plot_type == "line chart":
            df.plot(kind="line")
        else:
            return f"暂不支持{plot_type}类型的图表"

        plt.title(f"{plot_type.title()} 查询结果")
        plt.tight_layout()
        plt.savefig("query_plot.png")
        plt.close()

        return f"已成功生成{plot_type.title()}，保存为query_plot.png"

    async def _arun(self, input_str: str, run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        raise NotImplementedError("sql_plot_tool 不支持异步调用")

步骤2：整合工具到SQL Agent #

将自定义工具与SQLDatabaseToolkit的工具合并，创建支持可视化的Agent：

# 获取SQL工具集
sql_tools = toolkit.get_tools()
# 初始化自定义绘图工具
plot_tool = SQLPlotTool(db=db)
# 合并所有工具
all_tools = sql_tools + [plot_tool]

# 创建Agent
agent_executor = create_sql_agent(
    llm=llm,
    tools=all_tools,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    input_variables=["question", "agent_scratchpad"],
    suffix=suffix,
    verbose=True,
    handle_parsing_errors=True
)

方法2：整合Python REPL工具（灵活但需注意安全） #

允许LLM先调用SQL工具获取数据，再使用Python REPL执行绘图代码。这种方式更灵活，但需注意REPL执行代码的安全性。

步骤1：添加Python REPL工具 #

from langchain.tools import PythonREPLTool

# 获取SQL工具集
sql_tools = toolkit.get_tools()
# 添加Python REPL工具
python_repl = PythonREPLTool()
all_tools = sql_tools + [python_repl]

步骤2：调整Agent提示词 #

修改提示词，引导LLM正确组合使用两个工具：

# 获取默认Agent提示词并修改
base_prompt = agent_executor.prompt.template
custom_prompt = base_prompt.replace(
    "You have access to the following tools:",
    """You have access to the following tools:
- SQL工具：用于查询Postgres数据库获取数据，返回结果可以转换为Pandas DataFrame
- Python REPL：用于执行Python代码，可使用matplotlib、seaborn等库绘制图表。步骤建议：
  1. 使用SQL工具查询数据
  2. 将查询结果转换为DataFrame
  3. 编写Python代码调用绘图库生成可视化"""
)

# 创建更新后的Agent
agent_executor = create_sql_agent(
    llm=llm,
    tools=all_tools,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    input_variables=["question", "agent_scratchpad"],
    suffix=suffix,
    verbose=True,
    handle_parsing_errors=True,
    prompt=custom_prompt
)

注意事项 #

• 方法1更适合生产环境，避免LLM执行未知Python代码的风险，且逻辑更可控；
• 方法2适合开发/测试场景，能支持更多自定义可视化需求，但需限制REPL的执行权限（比如禁止文件写入、网络请求等）；
• 如果在Jupyter Notebook环境中，可将plt.savefig替换为plt.show()直接展示图表；
• 若需要前端展示，可将生成的图片转为Base64编码后返回，或存储到静态文件服务器。

内容的提问来源于stack exchange，提问作者rakibul-mahin