AI 文摘

使用检索增强生成(RAG)增强SQL代理





作者: 大数据杂货铺 来源: 大数据杂货铺

在上一篇文章中,我们深入探讨了构造 SQL 代理的过程,以帮助我们通过查询数据库中的数据来回答问题。在本文中,我们将探讨如何通过合并高级分析功能来增强 SQL 代理的功能。想象一下,代理不仅能够提供基本的统计数字,例如客户支付的平均金额,而且还能够提供更高级和更有趣的见解。这包括识别数据库中用户或产品之间的相似性或确定经常取消其成员资格的用户的路由路径等任务。让我们讨论一下如何实现这些高级功能。

Teradata的高级分析功能

与其他数据库不同,Teradata 通过提供大量高级分析功能而脱颖而出,从数据清理和数据探索到模型训练、文本分析以及路径和模式分析功能。

其显著特点是所有这些功能都可以在数据库内无缝运行,无需设置单独的环境。执行这些函数时,它们将直接在数据库中进行处理,从而确保高性能。

例如,考虑数据库中的两个表:UserHistory和UserHistoryReferences 。使用该TD_VectorDistance函数,您可以在这些表之间找到相似的用户。查询语法如下:


SELECT target_id, reference_id, distancetype, CAST(distance AS DECIMAL(36,8)) AS distance          
FROM TD_VECTORDISTANCE (          
    ON target_mobile_data_dense AS TargetTable          
    ON ref_mobile_data_dense AS ReferenceTable DIMENSION          
    USING          
        TargetIDColumn('userid')          
        TargetFeatureColumns('CallDuration','DataCounter','SMS')          
        RefIDColumn('userid')          
        RefFeatureColumns('CallDuration','DataCounter','SMS')          
        DistanceMeasure('cosine')          
        TopK(2)          
) AS dt ORDER BY 3,1,2,4;

以下是来自 DB 的结果:


Target_ID    Reference_ID   DistanceType   Distance          
---------    -------------  -------------  --------          
1            5              cosine         0.45486518          
1            7              cosine         0.32604815          
2            5              cosine         0.02608923          
2            7              cosine         0.00797609          
3            5              cosine         0.02415054          
3            7              cosine         0.00337338          
4            5              cosine         0.43822243          
4            7              cosine         0.31184844

有关 Teradata 的高级分析详细信息,请查看此处的文档。

检索增强生成 (RAG)

为了便于您的代理了解如何使用这些功能,我建议采用一种称为检索增强生成 (RAG) 的技术。

此方法有助于根据查询查找相关说明。例如,如果我要求我的代理帮助我根据表UserHistory和 UserHistoryReferences,RAG 将有效地返回与此请求相关的适当语法和示例。

语法说明

为了获得 SQL 代理的最佳性能,语法指令应包含两条基本信息。首先,包括语法和每个参数的解释。其次,也是最重要的,提供例子。您提供的示例越多,代理生成的 SQL 语法就越准确。

让我们构建 RAG

要创建 RAG 系统,首先要准备文档。将这些文档转换为向量并将它们保存在向量数据库中,我们将其称为向量数据库。在此示例中,我将使用名为 FAISS 的 Vector DB。


# Import require lib          
from langchain.embeddings import OpenAIEmbeddings          
from langchain.vectorstores import FAISS

首先查看 Teradata 提供的文档。首先,准备包含解释和示例的语法说明。


syntax_1 = """          
Syntax Description :           
    TD_VectorDistance () ...          
  

  

          
Example :          
    TD_VectorDistance ( ... )          
"""          
  

  

          
syntax_2 = """          
...          
"""          
  

  

          
syntax_3 = """          
...          
"""

接下来,利用来自 Hugging Face 或 OpenAI Embedding Service 等平台的各种开源模型。在这种情况下,我利用 OpenAI 完成了这项任务。

embedding_function = OpenAIEmbeddings(openai_api_key=os.getenv("OPENAI_API"))

最后,在 Langchain 和 FAISS 数据库的帮助下,您只需几行代码即可完成该过程


technical_list = [syntax_1, syntax_2, syntax_3, ..syntax_n]         
          
db = FAISS.from_texts(technical_list, embedding_function)

您可以使用下面的简单代码轻松搜索数据库中的相关文档。例如,如果要查找用于计算相似性的语法,则以下代码将返回在上一步中准备的与查询相关的确切语法:

db.similarity_search("Calculate similarity")[0]

将 RAG 与 SQL 代理集成

我们已经在之前的博客文章中介绍了如何创建 SQL 代理。如果您不熟悉该过程,请参阅该博客。此外,我们还讨论了创建一个 RAG 来检索相关的语法指令信息。现在,让我们探讨如何无缝集成这两个组件。

RAG 作为工具

在这篇博客中,我详细介绍了 如何利用SQL 代理与数据库交互等工具。现在,我的想法是将 RAG 指定为另一个工具sql_db_list_tables。这允许 SQL 代理决定何时浏览相关文档,并在需要时确定最合适的关键字进行搜索。

创建自定义工具

要使用 Langchain 创建自定义工具,请扩展 Langchain 提供的类并自定义函数BaseTool_run,如下所示。保持说明清晰至关重要,以确保 SQL 代理了解该工具的用途。


from langchain.tools import BaseTool          
from typing import Optional          
from langchain.callback_manager import CallbackManagerForToolRun, AsyncCallbackManagerForToolRun          
  

  

          
# Define retriever          
retriever = db.as_retriever()          
  

  

          
# Define customize tool          
class TeradataSearchTool(BaseTool):          
    name = "teradata_search_tool"          
    description = "Input to this tool is a keyword such as binning or bucketing, similarity, moving average. Output is an instruction on how to use Teradata Syntax with examples to improve queries."          
  

  

          
    def _run(          
        self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None          
    ) -> str:          
        """Use the tool."""          
        global retriever          
        relevant_doc = retriever.get_relevant_documents(query)          
        if len(relevant_doc) == 0 or len(query) == 0:          
            return "There are no Teradata syntax examples to be used in this scenario."          
        else:          
            return relevant_doc[0].page_content          
  

  

          
    async def _arun(          
        self, query: str, run_manager: Optional[AsyncCallbackManagerForToolRun] = None          
    ) -> str:          
        """Use the tool asynchronously."""          
        raise NotImplementedError("custom_search does not support async")          
  

  

          
# Init teradata search tool          
teradata_search_tool = TeradataSearchTool()

####使用自定义工具创建 SQL 代理

定义 Teradata 搜索工具后,可以使用以下代码创建 SQL 代理。此外,您可以在“extra_tools”部分中包括在上一步中创建的“teradata_search_tool”。


# Step 4. Create Agent Executor           
sql_agent = create_sql_agent(          
    llm=model,          
    toolkit=toolkit,          
    verbose=True,          
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,          
    extra_tools=[teradata_search_tool],          
    prefix=prefix,          
    suffix=suffix          
)

最后,测试一下

agent_executor.run("Identify user similarities by analyzing the 'UserHistory' table using 'UserHistoryReference' as the reference table, focusing on attributes CallDuration, DataCounter, and SMS")

当我要求代理通过使用“UserHistoryReference”作为参考表来分析“UserHistory”表来帮助我识别用户相似性时,重点关注 CallDuration、DataCounter 和 SMS 等属性,以下是结果。

代理结果

结论

通过将 SQL 代理与 RAG 相结合,我们将 LLM 模型的强大功能提升到一个新的水平。此方法允许创建另一个 RAG,使您的代理能够根据结构化数据和文本数据回答问题。但是,必须承认仍然存在一些与令牌限制相关的问题。在下一篇博客中,我将深入讨论如何微调模型,使您的代理能够在不依赖 RAG 的情况下执行相同的任务。

原文链接:https://medium.com/@lucnguyen_61589/enhancing-sql-agents-with-retrieval-augmented-generation-rag-e20dbd8bb685

更多AI工具,参考Github-AiBard123国内AiBard123

可关注我们的公众号:每天AI新工具