Vizro-AI与LangChain集成指南：实现智能可视化工具链

引言

在当今数据驱动的决策环境中，将自然语言处理与数据可视化相结合已成为提升分析效率的关键。Vizro-AI作为McKinsey开源的可视化工具，通过与LangChain框架的集成，为用户提供了更智能、更灵活的数据可视化解决方案。本文将详细介绍如何将Vizro-AI的功能封装为LangChain工具，构建一个完整的可视化工具链。

环境准备

在开始集成前，我们需要确保环境配置正确。以下是基础环境设置代码：

from copy import deepcopy
from typing import Annotated, Any

import pandas as pd
import vizro.plotly.express as px
from langchain_core.runnables import chain
from langchain_core.tools import InjectedToolArg, tool
from langchain_openai import ChatOpenAI
from vizro_ai import VizroAI

# 初始化语言模型，推荐使用GPT-4以获得最佳效果
llm = ChatOpenAI(model="gpt-4")

工具定义与封装

核心设计理念

Vizro-AI与传统LangChain工具的主要区别在于它需要处理Pandas DataFrame作为输入。直接传递完整数据集不仅效率低下，也存在安全隐患。我们采用LangChain的运行时绑定特性来解决这一问题。

图表生成工具

@tool(parse_docstring=True)
def get_plot_code(df: Annotated[Any, InjectedToolArg], question: str) -> str:
    """生成可视化图表代码
    
    参数:
        df: Pandas DataFrame数据
        question: 描述图表需求的自然语言
        
    返回:
        生成的Vizro图表代码
    """
    vizro_ai = VizroAI(model=llm)
    plot_elements = vizro_ai.plot(
        df,
        user_input=question,
        return_elements=True,
    )
    return plot_elements.code_vizro

仪表盘生成工具

@tool(parse_docstring=True)
def get_dashboard_code(dfs: Annotated[Any, InjectedToolArg], question: str) -> str:
    """生成完整仪表盘代码
    
    参数:
        dfs: 多个Pandas DataFrame组成的列表
        question: 描述仪表盘需求的自然语言
        
    返回:
        生成的Vizro仪表盘代码
    """
    vizro_ai = VizroAI(model=llm)
    dashboard_elements = vizro_ai.dashboard(
        dfs,
        user_input=question,
        return_elements=True,
    )
    return dashboard_elements.code

工具链构建

数据注入机制

# 将工具绑定到语言模型
tools = [get_plot_code, get_dashboard_code]
llm_with_tools = llm.bind_tools(tools)

# 创建数据注入链
@chain
def inject_df(ai_msg):
    tool_calls = []
    for tool_call in ai_msg.tool_calls:
        tool_call_copy = deepcopy(tool_call)

        if tool_call_copy["name"] == "get_dashboard_code":
            tool_call_copy["args"]["dfs"] = dfs
        else:
            tool_call_copy["args"]["df"] = df

        tool_calls.append(tool_call_copy)
    return tool_calls

工具路由与执行

# 创建工具映射
tool_map = {tool.name: tool for tool in tools}

@chain
def tool_router(tool_call):
    return tool_map[tool_call["name"]]

# 组合完整工具链
chain = llm_with_tools | inject_df | tool_router.map()

实际应用示例

单图表生成

# 加载示例数据
df = px.data.gapminder()

# 生成各大洲人均GDP图表
plot_response = chain.invoke("Plot GDP per capita for each continent")
print(plot_response[0].content)

生成的代码示例：

import plotly.graph_objects as go
from vizro.models.types import capture

@capture("graph")
def custom_chart(data_frame):
    continent_gdp = data_frame.groupby("continent")["gdpPercap"].mean().reset_index()
    fig = go.Figure(data=[go.Bar(x=continent_gdp["continent"], y=continent_gdp["gdpPercap"])])
    fig.update_layout(
        title="GDP per Capita by Continent",
        xaxis_title="Continent",
        yaxis_title="GDP per Capita",
    )
    return fig

仪表盘生成

dfs = [px.data.gapminder()]

dashboard_response = chain.invoke(
    "Create a dashboard. This dashboard has a chart showing the correlation between gdpPercap and lifeExp."
)
print(dashboard_response[0].content)

生成的仪表盘代码包含完整的页面布局、图表组件和样式配置。

最佳实践建议

1. 数据预处理：在将数据传递给工具链前，确保数据已进行必要的清洗和转换
2. 提示工程：精心设计自然语言提示可以获得更符合需求的输出
3. 错误处理：为工具链添加适当的错误处理机制，特别是处理大型数据集时
4. 性能优化：对于复杂可视化需求，考虑分步执行或增加约束条件

结语

通过将Vizro-AI与LangChain集成，我们构建了一个强大的可视化工具链，能够将自然语言指令转化为专业的数据可视化。这种集成不仅提高了数据分析的效率，也使可视化工具的易用性达到了新的水平。开发者可以根据实际需求扩展这一基础框架，添加更多定制化功能。