不写SQL就能构建数据分析AI代理?

借助LangChain和DuckDB从零开始构建全面的AI代理

SQL一直是大多数数据分析任务的基础语言. 我们经常提出需要数据来解答的问题. SQL对于将业务需求转换为可执行的数据检索代码至关重要. 然而, 借助 AI , 我们可以开发出无需SQL专业知识即可处理业务查询的 AI 代理.

本文将演示如何使用LangChain和DuckDB从零开始构建 AI 代理. 在构建好自己的SQL AI 代理后, 您可以快速完成数据分析任务. 更令人感兴趣的是, 我们可以使用Kaggle上的随机数据集来测试 AI 在SQL中的数据分析能力.

今日使用的工具与数据

• DuckDB: DuckDB 是一个内存中 SQL OLAP 数据库管理系统；它易于安装, 并能从 CSV/JSON 文件中快速编写分析 SQL 查询.
• LangChain: LangChain是一个可组合的框架, 用于构建基于大型语言模型（LLM）的应用. 鉴于AI代理的快速发展, 开发代理的关键组件之一是使用LangChain等框架来管理LLM流程并执行特定的内置操作.
• OpenAI API密钥 : LangChain仅是一个用于构建基于LLM应用的框架. 为了理解分析问题, 使用已知模式解析它并生成结果, 我们仍然需要一个 LLM 解决方案. OpenAI 提供的强大 LLM 之一是⁣,gpt-4o, 这是 2024 年 5 月开发的最新模型. 我们将使用此模型作为我们 SQL 生成 AI 代理的"大脑".
• Netflix数据集 (CC0: 公共领域) : 我们将使用 Kaggle 上的 Netflix 电影和电视剧 数据集作为示例, 测试我们的 SQL 生成 AI 代理.

SQL生成AI代理的高级工作流程

SQL生成AI代理工作流程

从整体来看, 我们的SQL生成AI代理遵循三个主要步骤.

1. write_query: 所有核心逻辑在此步骤实现, 输入数据的理解是生成SQL的关键. 我们提供一个分析问题作为输入, 并获得一个SQL查询作为输出. 例如, 我们可以输入以下问题: "请给我年总收入超过两百万的商店." 输出将是一个可执行的SQL查询, 如:

SELECT store, SUM(sale_price) revenue
FROM sales
GROUP BY store
HAVING SUM(sale_price) > 2000000000

2. execute_query: 在从write_query步骤获取SQL语句后, 该步骤使用LLM生成的SQL在DuckDB上执行以获取实际结果. 这与手动编写并执行SQL语句的过程本质上相同. 我们可以利用LangChain管理此工作流的状态: write_query的输出将作为输入传递给execute_query.

3. generate_answer: 尽管步骤 2 的结果以表格格式获取, 但仍可能难以理解. 在 generate_answer 步骤中, 我们将表格结果重新输入 LLM, 期望其生成更易于人类阅读的输出.

通过遵循这三个步骤, 我们将能够创建一个可靠的 AI 代理, 无需了解 SQL 语法即可协助处理业务查询. 以下是上述三个步骤的输出摘要.

我提出的问题是: "你能获取每位导演的总演出场次, 并按总场次降序排列前 3 名导演吗?"

{'write_query': {'query': "SELECT director, COUNT(*) as total_shows \nFROM read_csv_auto('data/netflix_titles.csv') \nWHERE director IS NOT NULL \nGROUP BY director \nORDER BY total_shows DESC \nLIMIT 3;"}}
{'execute_query': {'result': [Document(metadata={}, page_content='director: Rajiv Chilaka\ntotal_shows: 19'), Document(metadata={}, page_content='director: Raúl Campos, Jan Suter\ntotal_shows: 18'), Document(metadata={}, page_content='director: Suhas Kadav\ntotal_shows: 16')]}}
{'generate_answer': {'answer': 'The total shows per director, sorted by total shows in descending order for the top 3 directors, are as follows:\n\n1. Rajiv Chilaka - 19 shows\n2. Raúl Campos, Jan Suter - 18 shows\n3. Suhas Kadav - 16 shows'}}

步骤 0. 选择适合 SQL 的 LLM.

首先, 我们需要正确配置所有环境变量, 选择gpt-4o作为基础LLM模型, 然后利用LangChain hub中的sql-query-system-prompt来简化SQL生成过程.

import os
from langchain_openai import ChatOpenAI
from langchain import hub

## Setup Proper Enviroment Variables 
os.environ["LANGCHAIN_API_KEY"] = os.environ.get("LANGCHAIN_API_KEY")
os.environ["LANGCHAIN_TRACING_V2"] = os.environ.get("LANGCHAIN_TRACING_V2")
os.environ["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY")

## Set the LLM Model
llm = ChatOpenAI(model="gpt-4o")

## Choose specific prompt to help with the model
query_prompt_template = hub.pull("langchain-ai/sql-query-system-prompt")

sql-query-system-prompt是一个提示模板, 后面跟有4个关键参数需要填写: dialect, top_k, table_info和input. 这些参数将在下一步中填写.

步骤1. 构建write_query函数

为了使工作流的输入和输出在 LangChain 中可交换, 我们将创建一个名为 State 的类来帮助保存每个步骤的输出.

从步骤 0 开始, 我们有 sql-query-system-prompt 作为提示模板. 为了提高生成的 SQL 的精度, 有必要告知 LLM 我们感兴趣的特定 SQL 方言. 由于我们将使用 duckDB, 可以通过提供 duckdb 作为方言来实现.

另一个关键参数是 table_info, 它可为 LLM 提供表的模式或样本行等更多信息. 由于我们直接从 CSV 文件读取表, 将利用 duckDB 中的 read_csv_auto 函数并传入 CSV 文件路径. table_info 将作为表名, 并在最终的 SQL 语句中用作 from 子句.

最后, 我们将生成的 SQL 作为查询返回, 以便在下一步中对 duckDB 执行.

from typing_extensions import Annotated

class State():
    question: str
    query: str
    result: str
    answer: str

class QueryOutput(TypedDict):
    query: Annotated[str, ..., "Syntactically valid SQL query."]


def write_query(state: State):
    """Generate SQL query to fetch information."""
    prompt = query_prompt_template.invoke(
        {
            "dialect": "duckdb",
            "top_k": 10,
            "table_info": f"read_csv_auto('{file_path}')",
            "input": state["question"],
        }
    )
    structured_llm = llm.with_structured_output(QueryOutput)
    result = structured_llm.invoke(prompt)
    return {"query": result["query"]}

让我们先测试输出, 看看生成的 SQL 是否符合预期.

sql_query = write_query({"question": "Can you get the total shows per director, and sort by total shows in desending order?"})
print(sql_query)

输出是一个有效的 SQL 语句, 似乎能够有效地利用列.

{'query': "SELECT director, COUNT(*) as total_shows \nFROM read_csv_auto('data/netflix_titles.csv') \nWHERE director IS NOT NULL\nGROUP BY director \nORDER BY total_shows DESC \nLIMIT 10;"}

步骤 2. 构建 execute_query 函数

现在我们有了 SQL 语句, 可以通过 duckDB 运行它来进行测试. LangChain 提供的 duckDB 加载器 支持与 LangChain 的文档加载器接口进行集成.

可以从状态对象中获取输出并传递给DuckDBloader, 它将执行从LLM生成的查询并运行DuckDB, 同时继续执行步骤1中的SQL查询.

此步骤因多种原因易出错, 因为我们完全依赖LLM生成的SQL. 为防止数据丢失或未经授权的操作, 我们可选择不执行任何DDL类型的问题, 或引入人工审核流程. 此外, 可在数据库端实施权限控制, 以防止特定SQL查询的意外执行.

from langchain_community.document_loaders import DuckDBLoader
def execute_query(state: State):
    """Execute SQL query."""
    data = DuckDBLoader(state["query"]).load()
    return {'result': data}

在duckDB中执行的查询输出如下.

{'result': [
    Document(metadata={}, page_content='director: Rajiv Chilaka\ntotal_shows: 19'), 
    Document(metadata={}, page_content='director: Raúl Campos, Jan Suter\ntotal_shows: 18'), 
    Document(metadata={}, page_content='director: Suhas Kadav\ntotal_shows: 16')
    ]
}

步骤3. 构建generate_answer函数

该流程的最后一步是利用SQL结果生成更易于人类理解的段落. 若您能够直接解析SQL结果, 此步骤可能并非必要. 然而, 利用LLM对输出进行总结仍是一个极具价值的步骤.

generate_answer 步骤是一个简单的过程, 涉及整合所有信息并让 LLM 总结结果.

def generate_answer(state: State):
    """Answer question using retrieved information as context."""
    prompt = (
        "Given the following user question, corresponding SQL query, "
        "and SQL result, answer the user question.\n\n"
        f'Question: {state["question"]}\n'
        f'SQL Query: {state["query"]}\n'
        f'SQL Result: {state["result"]}'
    )
    response = llm.invoke(prompt)
    return {"answer": response.content}

以下是 generate_answer 步骤的输出, LLM 还提供了换行符以方便作为段落阅读.

{'answer': 
'The total shows per director, sorted by total shows in descending order for the top 3 directors, are as follows:\n\n
  1. Rajiv Chilaka - 19 shows\n
  2. Raúl Campos, Jan Suter - 18 shows\n
  3. Suhas Kadav - 16 shows'
}

整合所有步骤

最终, 我们可以利用 LangChain 构建一个图, 从而实现工作流管理和状态控制的无缝自动化.

from langgraph.graph import START, StateGraph

graph_builder = StateGraph(State).add_sequence(
    [write_query, execute_query, generate_answer]
)

graph_builder.add_edge(START, "write_query")
graph = graph_builder.compile()

使用 LangChain Smith 监控请求

为了更深入地了解整个请求生命周期, 我们可以使用 LangChain Smith 获取更详细的信息. 该工具有效地为用户提供了对发送到我们 AI 代理的请求的可见性, 防止我们将 LLM 视为一个简单的黑盒.

以下是LangChain Smith提供的LangGraph. 我们可以获取图中每个步骤的详细信息, 以及每个步骤相关的延迟, 输入和输出.

LangChain Smith用于示例查询

附加示例

让我们使用AI代理对Netflix数据集提出更多问题, 以验证其是否能生成准确的SQL并提供正确答案.

Q1: Can you get the number of shows that start with letter D?

该查询准确识别了相应的列标题, 并正确统计了节目数量.

for step in graph.stream(
    {"question": "Can you get the number of shows that start with letter D??"}, stream_mode="updates"
):
    print(step)


""" Output is 
{'write_query': {'query': "SELECT COUNT(*) FROM read_csv_auto('data/netflix_titles.csv') WHERE title LIKE 'D%';"}}
{'execute_query': {'result': [Document(metadata={}, page_content='count_star(): 375')]}}
{'generate_answer': {'answer': 'There are 375 shows that start with the letter D.'}}
"""

Q2: Can you get the how many years between each show director Rajiv Chilaka produced, sort by release years?

这一点令我惊讶. 随着我们使用窗口函数构建越来越复杂的SQL, LLM能够有效理解该查询.

for step in graph.stream(
    {"question": "Can you get the how many years between each show director Rajiv Chilaka produced, sort by release years?"}, stream_mode="updates"
):
    print(step)

""" Output is
{'write_query': {'query': "SELECT title, release_year, release_year - LAG(release_year) OVER (ORDER BY release_year) AS years_between_releases\nFROM read_csv_auto('data/netflix_titles.csv')\nWHERE director = 'Rajiv Chilaka'\nORDER BY release_year\nLIMIT 10;"}}
{'execute_query': {'result': [Document(metadata={}, page_content='title: Chhota Bheem & Ganesh\nrelease_year: 2009\nyears_between_releases: None'), Document(metadata={}, page_content='title: Chhota Bheem aur Krishna\nrelease_year: 2009\nyears_between_releases: 0'), Document(metadata={}, page_content='title: Chhota Bheem & Krishna: Pataliputra- City of the Dead\nrelease_year: 2010\nyears_between_releases: 1'), Document(metadata={}, page_content='title: Chhota Bheem: Bheem vs Aliens\nrelease_year: 2010\nyears_between_releases: 0'), Document(metadata={}, page_content='title: Chhota Bheem & Krishna: Mayanagari\nrelease_year: 2011\nyears_between_releases: 1'), Document(metadata={}, page_content='title: Chhota Bheem: Journey to Petra\nrelease_year: 2011\nyears_between_releases: 0'), Document(metadata={}, page_content='title: Chhota Bheem: Master of Shaolin\nrelease_year: 2011\nyears_between_releases: 0'), Document(metadata={}, page_content='title: Chhota Bheem Aur Hanuman\nrelease_year: 2012\nyears_between_releases: 1'), Document(metadata={}, page_content='title: Chhota Bheem: Dholakpur to Kathmandu\nrelease_year: 2012\nyears_between_releases: 0'), Document(metadata={}, page_content='title: Chhota Bheem: The Rise of Kirmada\nrelease_year: 2012\nyears_between_releases: 0')]}}
{'generate_answer': {'answer': 'Based on the SQL result, here are the number of years between each show that director Rajiv Chilaka produced, sorted by release years:\n\n1. "Chhota Bheem & Ganesh" - Released in 2009 (No previous release to compare, so the difference is None)\n2. "Chhota Bheem aur Krishna" - Released in 2009 (0 years since the previous release)\n3. "Chhota Bheem & Krishna: Pataliputra- City of the Dead" - Released in 2010 (1 year since the previous release)\n4. "Chhota Bheem: Bheem vs Aliens" - Released in 2010 (0 years since the previous release)\n5. "Chhota Bheem & Krishna: Mayanagari" - Released in 2011 (1 year since the previous release)\n6. "Chhota Bheem: Journey to Petra" - Released in 2011 (0 years since the previous release)\n7. "Chhota Bheem: Master of Shaolin" - Released in 2011 (0 years since the previous release)\n8. "Chhota Bheem Aur Hanuman" - Released in 2012 (1 year since the previous release)\n9. "Chhota Bheem: Dholakpur to Kathmandu" - Released in 2012 (0 years since the previous release)\n10. "Chhota Bheem: The Rise of Kirmada" - Released in 2012 (0 years since the previous release)'}}
"""

最后思考

通过上述示例, 我们可以制定复杂的查询以从LLM获取响应, 并高效地从duckDB数据库中检索请求. 由于LLM技术的进步, 进行数据分析不再需要具备SQL背景. 这可能引发关于数据分析师角色的新问题: LLM生成的SQL是否会提升他们的能力, 还是会威胁到他们的职位?

本文观察到AI代理能够无障碍完成特定任务. 尽管AI代理仍需人类干预的领域存在, 但人类所需参与的时间比例已显著降低. 我预计到2025年, 针对特定领域问题的AI代理构建将呈现增长趋势.

好吧, 今天的内容就分享到这里啦!

一家之言, 欢迎拍砖!

Happy coding! Stay GOLDEN!