提升sql生成能力：利用上下文优化agentic Sql生成的3种方法

Rifx.Online
Large Language Models , Data Science , AI Applications
23 Feb, 2025

当大型语言模型（LLMs）和生成性人工智能首次成为广泛流行的概念时

我认为这些模型将完全消除从头编写 SQL 的需求。我假设 SQL 的高度结构化特性将减少这些模型需要解决的不确定性，从而使可靠预测下一行输出变得更容易。显然，我的想法是错误的，因为我过于关注 SQL 查询的子句，而不是它的 上下文。是的，SQL 的高度结构化语法使得人类和模型都更容易学习。但编写 SQL 还需要深入理解查询运行的上下文，比如知道使用哪些表和列，理解连接关系等等。

有一些产品，如 Gemini 在 BigQuery 中或 Microsoft Copilot，能够将一些关键上下文与您的提示一起传递给模型，以改善 SQL 生成。然而，这些产品通常专注于直接在控制台中编写 SQL 的技术从业者，这些 SQL 旨在立即运行。这对于一些数据分析工作负载是有效的，比如数据分析师进行的临时分析。但对于那些在控制台外工作、不自己编写 SQL 的从业者，或者当您与不需要主键强制执行的数据仓库/数据库一起工作时，又该如何呢？例如：

正在构建代理或代理工作流的软件工程师，这些工作流使用基础模型如 Gemini 模型来生成 SQL
为 ELT 工作流编写查询的数据工程团队
在集成开发环境中开发模型的数据科学家

在这篇文章中，我们将看到一些如何使用 BigQuery 生成重要元数据并将其提供给基础模型以改善生成的 SQL 的示例。这些示例主要旨在支持开发生成 SQL 的代理工作流的从业者，这些工作流作为其工作流的一部分使用 Gemini 基础模型，但我也会包括一些如何支持其他用例的示例。我们将使用 TheLook 电子商务数据集，这是一个合成的公共数据集，代表一个虚构的电子商务零售商，用于本博客文章。

目标

我们的目标是使用生成性人工智能来确定位于日本的哪些用户在 2023 年下了最多的订单。这将要求我们联接两个表（用户和订单），以便我们可以返回用户的电子邮件和他们在 2023 年下的订单数量。让我们开始吧！

元数据生成

表结构

提高大型语言模型生成的SQL质量和可用性的第一步是定义模型可用数据的结构。当您在BigQuery中使用Gemini生成SQL时，此上下文会自动包含，但对于生成自己SQL的代理工作流，我们需要手动提供该信息。第一步是解析相关BigQuery表和数据集中的元数据，如列名和描述。我们可以使用这个Python函数来利用BigQuery Python 客户端库遍历每个数据集和表，然后将每个表的元数据写入JSON文件，以便我们稍后可以将其与提示一起传递给模型。

该函数还记录了表和数据集模式更新时间的时间戳，并将其写入pickle 文件。这使我们能够比较上次本地更新表元数据的时间，以查看是否需要刷新它。我们还可以通过BigQuery的表抽样进一步提升这一点。我们可以使用表抽样将每个表的一个子集包含在元数据JSON文件中，以便模型能够更好地理解我们的环境。请记住，这可能会带来数据安全风险，因此请确保您完全理解实施此类内容的影响。

连接关系

描述每个单独的表是一个好的开始，但对于使用多个表的查询，我们还需要模型理解它们应该如何连接。BigQuery 确实支持主键和外键的强制执行，这会告诉模型如何连接我们刚刚创建的模式文件中的表，但它并不要求你的模式包含它们。如果你的表没有定义主键和外键，你需要在表之间定义这些连接关系，但手动执行这一操作非常耗费人力。相反，我们可以使用以下方式生成这些关系：

BigQuery 的 INFORMATION_SCHEMA，它包含许多系统定义的视图，提供有关你的 BigQuery 对象的元数据信息。我们将使用 JOBS 视图来获取 SQL 查询文本。
使用 Gemini Pro 解析查询文本，并根据在我们的项目中运行的查询识别连接关系。

我们将使用 BigQuery 的 INFORMATION_SCHEMA jobs 视图来识别具有连接关系的查询，并将查询文本连接成一个单一的数组。我们将把这些查询的文本传递给我们的提示给模型，并让 Gemini Pro 生成描述这些连接关系的实体关系图 (ERD)。模型的响应是每个连接关系的摘要，写入一个 markdown 文件和一个 JSON 文件，我们稍后将在提示模型生成 SQL 时将其作为上下文的一部分提供。以下是 JSON 文件的示例：

生成 SQL

代理工作流

现在我们已经生成了所需的上下文，可以让 Gemini Flash 生成 SQL 来回答一个问题！Gemini Flash 的性能和成本优势使其成为请求相对简单的用例的绝佳选择，并且它仍然提供一个巨大的 100 万令牌上下文窗口，以传入相关的表结构和 ERD。下面的代码处理读取元数据和 ERD 文件的过程，然后从我们的模板生成一个提示，以传递给 Gemini 模型。

我们可以将这个 Python 函数配置为工具的一部分，以供处理我们的代理工作流的 Gemini 模型使用，使该函数能够处理 SQL 生成，而不是期望模型理解它需要在 SQL 生成之前本地提取上下文。但是……这真的能改善结果吗？看看下面当我问模型“生成一个 SQL 查询，以确定 2023 年在日本下单最多的用户”时的结果。

SQL 由 Gemini Pro 生成，左侧是没有额外上下文的情况，右侧是有上下文的情况。没有上下文的查询是有效的 SQL 语法，但没有包括调用不同表中列所需的连接。带有上下文生成的查询不仅是有效的 SQL，而且特别针对我们的问题！

生成既准确又有效的 SQL 对于代理和代理工作流来说更加重要，因为没有人类参与来捕捉和排除由错误查询引起的错误。这意味着查询需要正确定义度量标准和连接关系，并且要表示有效的 SQL 方言语法。让我们再看看生成 SQL 的工作流。上面的函数使用 Gemini Flash 解析输入并生成初始 SQL 查询（第 20 行）。然后，它将模型生成的 SQL 以及原始输入（包括表结构和 ERD）传递给第二个 Gemini 模型调用，该调用专门检查查询以确保其具有有效的语法和表引用（第 25 行）。对于需要极其复杂 SQL 或在验证 SQL 时遇到问题的工作流，可以使用 Gemini Pro 模型来生成查询、验证查询或两者兼而有之！由于查询生成是可供模型使用的函数，作为代理工作流的一部分，它可以使用不同的模型版本，甚至是与用户在前端交互的模型完全不同的模型。例如：使用 Gemini Flash 的聊天机器人可以调用 Gemini Pro 来生成 SQL，并使用经过微调的 CodeGemma 模型来验证查询。您甚至可以结合可选的查询干运行来验证查询并返回错误消息，以反馈到第二个模型调用中以纠正问题，或者在运行之前估算查询的成本。

IDEs

Gemini Code Assist 在您喜欢的集成开发环境中提供代码生成和补全，例如 VS Code。此外，您可以快速与 Gemini Code Assist 聊天，以获取编码问题的答案或获得编码最佳实践的指导。对于我们的用例，这对不常编写 SQL 的软件工程师或正在开发自定义模型的数据科学家来说，将非常有用，因为他们可能希望使用 Gemini Code Assist 为他们生成 SQL 语句，或者节省编写 SQL 的时间。但是，就像基础模型一样，Code Assist 仍然需要额外的上下文来正确生成满足您用例的 SQL。当 Code Assist 没有这些上下文时，会发生以下情况，尽管我们告诉它使用哪个数据集和项目：

SQL 由 Gemini Code Assist 生成，但没有额外的上下文。这不仅不是有效的 BigQuery SQL，而且国家的过滤器使用了日本的 ISO 3166–2 代码 (JP)，而不是像用户表那样的完整文本名称。

这令人沮丧地接近！乍一看，这看起来像是有效的 BigQuery SQL 语法，尽管缺少 GROUP BY 语句。但是，即使查询验证器捕获了这个错误并且我们进行了更正，它仍然无法正确回答我们的问题，因为它假设用户表中的国家列包含每个国家的 2 字母 ISO 代码 (JP)，而实际上该列包含的是完整的文本名称（日本）。因此，这个查询将成功运行，但不会返回任何结果，因为没有行的国家等于 JP。让我们看看当我们传入 ERD 和表结构作为上下文时会发生什么：

SQL 由 Gemini Code Assist 生成，并提供了额外的上下文。这一个有效！

限制

正如我们所看到的，这种方法不仅可以显著提高生成性人工智能生成的查询质量，而且可以在不需要大量人力手动定义连接关系或确保表结构最新的情况下实现。但这并不是一种完美的方法，仍然存在一些限制。

跨项目连接

目前，这种方法仅包括定义项目内表的元数据。虽然没有理由不能扩展到多个项目，但这可能会导致下面定义的问题。

可扩展性

这种方法解析给定项目中的所有表，并在本地存储它们的元数据，以便在模型运行时可用。这是最直接的方法，但如果您的项目有大量表，或者如果您包含多个项目，它将遇到限制。这可以调整为将包含表列描述和数据样本的文件卸载到 Google Cloud Storage (GCS)，并仅在本地存储数据集和带有表描述的表名。这将允许模型识别所需的表，然后仅将相关的表元数据提取到本地内存中，但这可能需要为您的代理工作流创建额外的函数调用。

最终，这需要在本地存储每个表的列描述的存储成本与潜在的规模限制之间取得平衡，以及在 SQL 生成之前识别相关表并从远程位置下载元数据的额外工作流步骤的网络成本和性能权衡。一种潜在的方法是：使用 BigQuery 的 INFORMATION SCHEMA jobs view 来确定每个表的查询作业，然后在本地存储最受欢迎表的元数据，同时将其余部分卸载到 GCS。

连接关系

使用在项目中运行的查询文本来连接表是帮助记录可用于数据的现有业务上下文的好方法，但这并不是一个完美的方法。这里最大的限制是，查询历史中不存在的连接关系将不会在您的ERD中表示，因此生成查询的模型将无法意识到它们。这意味着：

尚未查询的表将不会定义其连接关系，因为模型将没有任何查询文本来解析和理解这些关系。
已经被查询的表可能在ERD中没有定义所有可能的连接关系，因为它仅考虑存在于查询中的连接。

此外，您的查询历史中可能存在低质量或不准确的连接关系。此示例中的 INFORMATION_SCHEMA 查询过滤查询文本，仅包括包含JOIN语句的查询作业，这些作业的状态为“DONE”，并且没有错误结果原因。但仅仅因为一个带有连接的查询成功运行，并不一定意味着它准确地定义了您数据中的有效关系。您可能需要添加额外的过滤器，仅包括由服务帐户运行的查询，这些服务帐户验证BI工作负载和ELT过程，或者仅包含最常见的连接，以确保捕获高质量、准确的连接关系。您还可以包括一个UNION语句，通过包括INFORMATION_SCHEMA中的视图定义来包括定义项目中视图的SQL。

下一步

查看 GitHub 上的代码，开始构建更好的方法来在您的代理和代理工作流中生成 SQL！

您是否有兴趣进一步探索？以下是我建议尝试的一些想法：

这种方法在您拥有包含详细、准确描述的表和列时将最为成功。但就像手动定义连接关系一样，这是一项令人烦恼的乏味任务。尝试使用这个笔记本，它使用 Gemini Pro 根据可用的元数据和列中的数据示例生成列和表描述。
解决一些限制。您可以从可扩展性开始，这需要您更改元数据输出位置，并创建一个新函数，以识别要查询的表并从 GCS 下载相关的元数据，然后再编写查询。

虽然这篇文章使用 Gemini 模型为 BigQuery 生成 GoogleSQL，但这种方法可以被推广/调整，以使其他生成性人工智能模型支持大多数数据仓库和数据库！数据库和数据仓库通常包含类似于 BigQuery 的信息架构，您可以使用它来定义连接关系和表结构。数据仓库通常提供查询历史，以便使用像我们在这里所做的 LLM 进行解析，您可以使用主键/外键为大多数数据库生成 ERD。这意味着您还可以使用这种方法生成 SQL，使用 Llama 模型进行验证，并在 Postgres 数据库中运行，或者使用其他模型解析和验证为 BigQuery 生成的 SQL。