构建基于 SQL 数据的问答系统

先决条件

本指南假定您熟悉以下概念

使 LLM 系统能够查询结构化数据在性质上可能与非结构化文本数据有所不同。在后者中，通常生成可以针对向量数据库搜索的文本，而结构化数据的方法通常是让 LLM 以 DSL（例如 SQL）编写和执行查询。在本指南中，我们将介绍在数据库中的表格数据上创建问答系统的基本方法。我们将介绍使用链和代理的实现。这些系统将使我们能够询问有关数据库中数据的问题，并获得自然语言答案。两者之间的主要区别在于，我们的代理可以循环查询数据库多次，直到回答问题为止。

⚠️ 安全注意事项 ⚠️

构建 SQL 数据库的问答系统需要执行模型生成的 SQL 查询。这样做存在内在风险。请确保您的数据库连接权限始终尽可能限定在您的链/代理的需求范围内。这将减轻但不能消除构建模型驱动系统的风险。有关一般安全最佳实践的更多信息，请参阅此处。

架构

从高层次来看，这些系统的步骤是

将问题转换为 SQL 查询：模型将用户输入转换为 SQL 查询。
执行 SQL 查询：执行查询。
回答问题：模型使用查询结果响应用户输入。

请注意，查询 CSV 中的数据可以遵循类似的方法。有关更多详细信息，请参阅我们关于 CSV 数据问答的操作指南。

设置

首先，获取所需的包并设置环境变量

%%capture --no-stderr
%pip install --upgrade --quiet langchain-community langchainhub langgraph

# Comment out the below to opt-out of using LangSmith in this notebook. Not required.
if not os.environ.get("LANGSMITH_API_KEY"):
    os.environ["LANGSMITH_API_KEY"] = getpass.getpass()
    os.environ["LANGSMITH_TRACING"] = "true"

示例数据

以下示例将使用带有 Chinook 数据库的 SQLite 连接，这是一个表示数字媒体商店的示例数据库。按照这些安装步骤在与此笔记本相同的目录中创建 Chinook.db。您还可以通过命令行下载和构建数据库

curl -s https://raw.githubusercontent.com/lerocha/chinook-database/master/ChinookDatabase/DataSources/Chinook_Sqlite.sql | sqlite3 Chinook.db

现在，Chinook.db 在我们的目录中，我们可以使用 SQLAlchemy 驱动的 SQLDatabase 类与之交互

from langchain_community.utilities import SQLDatabase

db = SQLDatabase.from_uri("sqlite:///Chinook.db")
print(db.dialect)
print(db.get_usable_table_names())
db.run("SELECT * FROM Artist LIMIT 10;")

API 参考：SQLDatabase

sqlite
['Album', 'Artist', 'Customer', 'Employee', 'Genre', 'Invoice', 'InvoiceLine', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']

"[(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains'), (6, 'Antônio Carlos Jobim'), (7, 'Apocalyptica'), (8, 'Audioslave'), (9, 'BackBeat'), (10, 'Billy Cobham')]"

太棒了！我们有了一个可以查询的 SQL 数据库。现在让我们尝试将其连接到 LLM。

链

链是可预测步骤的组合。在 LangGraph 中，我们可以通过简单的节点序列来表示链。让我们创建一个步骤序列，给定一个问题，执行以下操作

将问题转换为 SQL 查询；
执行查询；
使用结果回答原始问题。

有些情况此安排不支持。例如，此系统将为任何用户输入（甚至“hello”）执行 SQL 查询。重要的是，正如我们将在下面看到的，有些问题需要多次查询才能回答。我们将在“代理”部分中解决这些情况。

应用程序状态

LangGraph 状态控制应用程序的哪些数据被输入、步骤之间传输以及应用程序输出。它通常是 TypedDict，但也可以是 Pydantic BaseModel。

对于此应用程序，我们只需跟踪输入问题、生成的查询、查询结果和生成的答案

from typing_extensions import TypedDict


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

现在我们只需要对该状态进行操作并填充其内容的功能。

将问题转换为 SQL 查询

第一步是获取用户输入并将其转换为 SQL 查询。为了可靠地获取 SQL 查询（没有 markdown 格式和解释或说明），我们将使用 LangChain 的结构化输出抽象。

让我们为我们的应用程序选择一个聊天模型

选择聊天模型

pip install -qU "langchain[openai]"

import getpass
import os

if not os.environ.get("OPENAI_API_KEY"):
  os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter API key for OpenAI: ")

from langchain.chat_models import init_chat_model

llm = init_chat_model("gpt-4o-mini", model_provider="openai")

我们将从 Prompt Hub 中提取提示来指导模型。

from langchain import hub

query_prompt_template = hub.pull("langchain-ai/sql-query-system-prompt")

assert len(query_prompt_template.messages) == 1
query_prompt_template.messages[0].pretty_print()

API 参考：hub

================================[1m System Message [0m================================

Given an input question, create a syntactically correct [33;1m[1;3m{dialect}[0m query to run to help find the answer. Unless the user specifies in his question a specific number of examples they wish to obtain, always limit your query to at most [33;1m[1;3m{top_k}[0m results. You can order the results by a relevant column to return the most interesting examples in the database.

Never query for all the columns from a specific table, only ask for a the few relevant columns given the question.

Pay attention to use only the column names that you can see in the schema description. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.

Only use the following tables:
[33;1m[1;3m{table_info}[0m

Question: [33;1m[1;3m{input}[0m

提示包括我们需要填充的几个参数，例如 SQL 方言和表架构。LangChain 的 SQLDatabase 对象包含有助于实现此目的的方法。我们的 write_query 步骤将仅填充这些参数并提示模型生成 SQL 查询

from typing_extensions import Annotated


class QueryOutput(TypedDict):
    """Generated SQL query."""

    query: Annotated[str, ..., "Syntactically valid SQL query."]


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

让我们测试一下

write_query({"question": "How many Employees are there?"})

{'query': 'SELECT COUNT(EmployeeId) AS EmployeeCount FROM Employee;'}

执行查询

这是创建 SQL 链最危险的部分。 请仔细考虑是否可以对您的数据运行自动化查询。尽可能减少数据库连接权限。考虑在查询执行之前向您的链添加人工批准步骤（见下文）。

为了执行查询，我们将从 langchain-community 加载一个工具。我们的 execute_query 节点将仅包装此工具

from langchain_community.tools.sql_database.tool import QuerySQLDatabaseTool


def execute_query(state: State):
    """Execute SQL query."""
    execute_query_tool = QuerySQLDatabaseTool(db=db)
    return {"result": execute_query_tool.invoke(state["query"])}

API 参考：QuerySQLDatabaseTool

测试此步骤

execute_query({"query": "SELECT COUNT(EmployeeId) AS EmployeeCount FROM Employee;"})

{'result': '[(8,)]'}

生成答案

最后，我们的最后一步是根据从数据库中提取的信息生成问题的答案

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}

使用 LangGraph 编排

最后，我们将我们的应用程序编译成一个单独的 graph 对象。在本例中，我们只是将这三个步骤连接成一个序列。

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()

API 参考：StateGraph

LangGraph 还附带内置实用程序，用于可视化应用程序的控制流

from IPython.display import Image, display

display(Image(graph.get_graph().draw_mermaid_png()))

让我们测试一下我们的应用程序！请注意，我们可以流式传输各个步骤的结果

for step in graph.stream(
    {"question": "How many employees are there?"}, stream_mode="updates"
):
    print(step)

{'write_query': {'query': 'SELECT COUNT(EmployeeId) AS EmployeeCount FROM Employee;'}}
{'execute_query': {'result': '[(8,)]'}}
{'generate_answer': {'answer': 'There are 8 employees.'}}

查看 LangSmith 跟踪。

人工参与环节

LangGraph 支持许多对这种工作流程有用的功能。其中之一是人工参与环节：我们可以在敏感步骤（例如执行 SQL 查询）之前中断我们的应用程序，以供人工审查。这通过 LangGraph 的持久性层实现，该层将运行进度保存到您选择的存储中。下面，我们指定内存中的存储

from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()
graph = graph_builder.compile(checkpointer=memory, interrupt_before=["execute_query"])

# Now that we're using persistence, we need to specify a thread ID
# so that we can continue the run after review.
config = {"configurable": {"thread_id": "1"}}

API 参考：MemorySaver

display(Image(graph.get_graph().draw_mermaid_png()))

让我们重复相同的运行，添加一个简单的“是/否”批准步骤

for step in graph.stream(
    {"question": "How many employees are there?"},
    config,
    stream_mode="updates",
):
    print(step)

try:
    user_approval = input("Do you want to go to execute query? (yes/no): ")
except Exception:
    user_approval = "no"

if user_approval.lower() == "yes":
    # If approved, continue the graph execution
    for step in graph.stream(None, config, stream_mode="updates"):
        print(step)
else:
    print("Operation cancelled by user.")

{'write_query': {'query': 'SELECT COUNT(EmployeeId) AS EmployeeCount FROM Employee;'}}
{'__interrupt__': ()}
``````output
Do you want to go to execute query? (yes/no):  yes
``````output
{'execute_query': {'result': '[(8,)]'}}
{'generate_answer': {'answer': 'There are 8 employees.'}}

有关更多详细信息和示例，请参阅此 LangGraph 指南。

后续步骤

对于更复杂的查询生成，我们可能希望创建少样本提示或添加查询检查步骤。有关此类高级技术和更多信息，请查看

提示策略：高级提示工程技术。
查询检查：添加查询验证和错误处理。
大型数据库：用于处理大型数据库的技术。

代理

代理利用 LLM 的推理能力在执行期间做出决策。使用代理允许您将对查询生成和执行过程的额外自由裁量权转移出去。尽管他们的行为不如上述“链”那样可预测，但它们具有一些优势

他们可以根据需要多次查询数据库以回答用户问题。
他们可以通过运行生成的查询、捕获回溯并正确地重新生成它来从错误中恢复。
他们可以根据数据库的架构以及数据库的内容（如描述特定表）回答问题。

下面我们组装一个最小的 SQL 代理。我们将使用 LangChain 的 SQLDatabaseToolkit 为其配备一组工具。使用 LangGraph 的预构建的 ReAct 代理构造函数，我们可以在一行中完成此操作。

提示

查看 LangGraph 的 SQL 代理教程，了解更高级的 SQL 代理公式。

SQLDatabaseToolkit 包括可以执行以下操作的工具

创建和执行查询
检查查询语法
检索表描述
... 以及更多

from langchain_community.agent_toolkits import SQLDatabaseToolkit

toolkit = SQLDatabaseToolkit(db=db, llm=llm)

tools = toolkit.get_tools()

tools

API 参考：SQLDatabaseToolkit

[QuerySQLDatabaseTool(description="Input to this tool is a detailed and correct SQL query, output is a result from the database. If the query is not correct, an error message will be returned. If an error is returned, rewrite the query, check the query, and try again. If you encounter an issue with Unknown column 'xxxx' in 'field list', use sql_db_schema to query the correct table fields.", db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x10d5f9120>),
 InfoSQLDatabaseTool(description='Input to this tool is a comma-separated list of tables, output is the schema and sample rows for those tables. Be sure that the tables actually exist by calling sql_db_list_tables first! Example Input: table1, table2, table3', db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x10d5f9120>),
 ListSQLDatabaseTool(db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x10d5f9120>),
 QuerySQLCheckerTool(description='Use this tool to double check if your query is correct before executing it. Always use this tool before executing a query with sql_db_query!', db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x10d5f9120>, llm=ChatOpenAI(client=<openai.resources.chat.completions.Completions object at 0x119315480>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x119317550>, root_client=<openai.OpenAI object at 0x10d5f8df0>, root_async_client=<openai.AsyncOpenAI object at 0x1193154e0>, model_name='gpt-4o', temperature=0.0, model_kwargs={}, openai_api_key=SecretStr('**********')), llm_chain=LLMChain(verbose=False, prompt=PromptTemplate(input_variables=['dialect', 'query'], input_types={}, partial_variables={}, template='\n{query}\nDouble check the {dialect} query above for common mistakes, including:\n- Using NOT IN with NULL values\n- Using UNION when UNION ALL should have been used\n- Using BETWEEN for exclusive ranges\n- Data type mismatch in predicates\n- Properly quoting identifiers\n- Using the correct number of arguments for functions\n- Casting to the correct data type\n- Using the proper columns for joins\n\nIf there are any of the above mistakes, rewrite the query. If there are no mistakes, just reproduce the original query.\n\nOutput the final SQL query only.\n\nSQL Query: '), llm=ChatOpenAI(client=<openai.resources.chat.completions.Completions object at 0x119315480>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x119317550>, root_client=<openai.OpenAI object at 0x10d5f8df0>, root_async_client=<openai.AsyncOpenAI object at 0x1193154e0>, model_name='gpt-4o', temperature=0.0, model_kwargs={}, openai_api_key=SecretStr('**********')), output_parser=StrOutputParser(), llm_kwargs={}))]

系统提示

我们还需要为我们的代理加载系统提示。这将包括有关如何行为的说明。

from langchain import hub

prompt_template = hub.pull("langchain-ai/sql-agent-system-prompt")

assert len(prompt_template.messages) == 1
prompt_template.messages[0].pretty_print()

API 参考：hub

================================[1m System Message [0m================================

You are an agent designed to interact with a SQL database.
Given an input question, create a syntactically correct [33;1m[1;3m{dialect}[0m query to run, then look at the results of the query and return the answer.
Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most [33;1m[1;3m{top_k}[0m results.
You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for the relevant columns given the question.
You have access to tools for interacting with the database.
Only use the below tools. Only use the information returned by the below tools to construct your final answer.
You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.

DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.

To start you should ALWAYS look at the tables in the database to see what you can query.
Do NOT skip this step.
Then you should query the schema of the most relevant tables.

让我们填充提示中突出显示的参数

system_message = prompt_template.format(dialect="SQLite", top_k=5)

初始化代理

我们将使用预构建的 LangGraph 代理来构建我们的代理

from langchain_core.messages import HumanMessage
from langgraph.prebuilt import create_react_agent

agent_executor = create_react_agent(llm, tools, prompt=system_message)

API 参考：HumanMessage | create_react_agent

考虑代理如何响应以下问题

question = "Which country's customers spent the most?"

for step in agent_executor.stream(
    {"messages": [{"role": "user", "content": question}]},
    stream_mode="values",
):
    step["messages"][-1].pretty_print()

================================[1m Human Message [0m=================================

Which country's customers spent the most?
==================================[1m Ai Message [0m==================================
Tool Calls:
  sql_db_list_tables (call_tFp7HYD6sAAmCShgeqkVZH6Q)
 Call ID: call_tFp7HYD6sAAmCShgeqkVZH6Q
  Args:
=================================[1m Tool Message [0m=================================
Name: sql_db_list_tables

Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track
==================================[1m Ai Message [0m==================================
Tool Calls:
  sql_db_schema (call_KJZ1Jx6JazyDdJa0uH1UeiOz)
 Call ID: call_KJZ1Jx6JazyDdJa0uH1UeiOz
  Args:
    table_names: Customer, Invoice
=================================[1m Tool Message [0m=================================
Name: sql_db_schema


CREATE TABLE "Customer" (
	"CustomerId" INTEGER NOT NULL, 
	"FirstName" NVARCHAR(40) NOT NULL, 
	"LastName" NVARCHAR(20) NOT NULL, 
	"Company" NVARCHAR(80), 
	"Address" NVARCHAR(70), 
	"City" NVARCHAR(40), 
	"State" NVARCHAR(40), 
	"Country" NVARCHAR(40), 
	"PostalCode" NVARCHAR(10), 
	"Phone" NVARCHAR(24), 
	"Fax" NVARCHAR(24), 
	"Email" NVARCHAR(60) NOT NULL, 
	"SupportRepId" INTEGER, 
	PRIMARY KEY ("CustomerId"), 
	FOREIGN KEY("SupportRepId") REFERENCES "Employee" ("EmployeeId")
)

/*
3 rows from Customer table:
CustomerId	FirstName	LastName	Company	Address	City	State	Country	PostalCode	Phone	Fax	Email	SupportRepId
1	Luís	Gonçalves	Embraer - Empresa Brasileira de Aeronáutica S.A.	Av. Brigadeiro Faria Lima, 2170	São José dos Campos	SP	Brazil	12227-000	+55 (12) 3923-5555	+55 (12) 3923-5566	luisg@embraer.com.br	3
2	Leonie	Köhler	None	Theodor-Heuss-Straße 34	Stuttgart	None	Germany	70174	+49 0711 2842222	None	leonekohler@surfeu.de	5
3	François	Tremblay	None	1498 rue Bélanger	Montréal	QC	Canada	H2G 1A7	+1 (514) 721-4711	None	ftremblay@gmail.com	3
*/


CREATE TABLE "Invoice" (
	"InvoiceId" INTEGER NOT NULL, 
	"CustomerId" INTEGER NOT NULL, 
	"InvoiceDate" DATETIME NOT NULL, 
	"BillingAddress" NVARCHAR(70), 
	"BillingCity" NVARCHAR(40), 
	"BillingState" NVARCHAR(40), 
	"BillingCountry" NVARCHAR(40), 
	"BillingPostalCode" NVARCHAR(10), 
	"Total" NUMERIC(10, 2) NOT NULL, 
	PRIMARY KEY ("InvoiceId"), 
	FOREIGN KEY("CustomerId") REFERENCES "Customer" ("CustomerId")
)

/*
3 rows from Invoice table:
InvoiceId	CustomerId	InvoiceDate	BillingAddress	BillingCity	BillingState	BillingCountry	BillingPostalCode	Total
1	2	2021-01-01 00:00:00	Theodor-Heuss-Straße 34	Stuttgart	None	Germany	70174	1.98
2	4	2021-01-02 00:00:00	Ullevålsveien 14	Oslo	None	Norway	0171	3.96
3	8	2021-01-03 00:00:00	Grétrystraat 63	Brussels	None	Belgium	1000	5.94
*/
==================================[1m Ai Message [0m==================================
Tool Calls:
  sql_db_query_checker (call_AQuTGbgH63u4gPgyV723yrjX)
 Call ID: call_AQuTGbgH63u4gPgyV723yrjX
  Args:
    query: SELECT c.Country, SUM(i.Total) as TotalSpent FROM Customer c JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY c.Country ORDER BY TotalSpent DESC LIMIT 1;
=================================[1m Tool Message [0m=================================
Name: sql_db_query_checker

\`\`\`sql
SELECT c.Country, SUM(i.Total) as TotalSpent FROM Customer c JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY c.Country ORDER BY TotalSpent DESC LIMIT 1;
\`\`\`
==================================[1m Ai Message [0m==================================
Tool Calls:
  sql_db_query (call_B88EwU44nwwpQL5M9nlcemSU)
 Call ID: call_B88EwU44nwwpQL5M9nlcemSU
  Args:
    query: SELECT c.Country, SUM(i.Total) as TotalSpent FROM Customer c JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY c.Country ORDER BY TotalSpent DESC LIMIT 1;
=================================[1m Tool Message [0m=================================
Name: sql_db_query

[('USA', 523.06)]
==================================[1m Ai Message [0m==================================

The country whose customers spent the most is the USA, with a total spending of 523.06.

您还可以使用 LangSmith 跟踪来可视化这些步骤和关联的元数据。

请注意，代理执行多个查询，直到获得所需的信息

列出可用表；
检索三个表的架构；
通过连接操作查询多个表。

然后，代理能够使用最终查询的结果来生成原始问题的答案。

代理可以类似地处理定性问题

question = "Describe the playlisttrack table"

for step in agent_executor.stream(
    {"messages": [{"role": "user", "content": question}]},
    stream_mode="values",
):
    step["messages"][-1].pretty_print()

================================[1m Human Message [0m=================================

Describe the playlisttrack table
==================================[1m Ai Message [0m==================================
Tool Calls:
  sql_db_list_tables (call_fMF8eTmX5TJDJjc3Mhdg52TI)
 Call ID: call_fMF8eTmX5TJDJjc3Mhdg52TI
  Args:
=================================[1m Tool Message [0m=================================
Name: sql_db_list_tables

Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track
==================================[1m Ai Message [0m==================================
Tool Calls:
  sql_db_schema (call_W8Vkk4NEodkAAIg8nexAszUH)
 Call ID: call_W8Vkk4NEodkAAIg8nexAszUH
  Args:
    table_names: PlaylistTrack
=================================[1m Tool Message [0m=================================
Name: sql_db_schema


CREATE TABLE "PlaylistTrack" (
	"PlaylistId" INTEGER NOT NULL, 
	"TrackId" INTEGER NOT NULL, 
	PRIMARY KEY ("PlaylistId", "TrackId"), 
	FOREIGN KEY("TrackId") REFERENCES "Track" ("TrackId"), 
	FOREIGN KEY("PlaylistId") REFERENCES "Playlist" ("PlaylistId")
)

/*
3 rows from PlaylistTrack table:
PlaylistId	TrackId
1	3402
1	3389
1	3390
*/
==================================[1m Ai Message [0m==================================

The `PlaylistTrack` table is designed to associate tracks with playlists. It has the following structure:

- **PlaylistId**: An integer that serves as a foreign key referencing the `Playlist` table. It is part of the composite primary key.
- **TrackId**: An integer that serves as a foreign key referencing the `Track` table. It is also part of the composite primary key.

The primary key for this table is a composite key consisting of both `PlaylistId` and `TrackId`, ensuring that each track can be uniquely associated with a playlist. The table enforces referential integrity by linking to the `Track` and `Playlist` tables through foreign keys.

处理高基数列

为了过滤包含专有名词（如地址、歌曲名称或艺术家）的列，我们首先需要仔细检查拼写，以便正确过滤数据。

我们可以通过创建一个向量存储来实现这一点，该向量存储包含数据库中存在的所有不同的专有名词。然后，我们可以让代理在用户的问题中包含专有名词时，每次都查询该向量存储，以查找该词的正确拼写。这样，代理可以确保它在构建目标查询之前了解用户指的是哪个实体。

首先，我们需要我们想要的每个实体的唯一值，为此我们定义一个函数，将结果解析为元素列表

import ast
import re


def query_as_list(db, query):
    res = db.run(query)
    res = [el for sub in ast.literal_eval(res) for el in sub if el]
    res = [re.sub(r"\b\d+\b", "", string).strip() for string in res]
    return list(set(res))


artists = query_as_list(db, "SELECT Name FROM Artist")
albums = query_as_list(db, "SELECT Title FROM Album")
albums[:5]

['In Through The Out Door',
 'Transmission',
 'Battlestar Galactica (Classic), Season',
 'A Copland Celebration, Vol. I',
 'Quiet Songs']

使用此函数，我们可以创建一个检索器工具，代理可以自行决定执行该工具。

让我们为此步骤选择一个嵌入模型和向量存储

选择一个嵌入模型:

选择嵌入模型

pip install -qU langchain-openai

import getpass
import os

if not os.environ.get("OPENAI_API_KEY"):
  os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter API key for OpenAI: ")

from langchain_openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

选择一个向量存储:

选择向量存储

pip install -qU langchain-core

from langchain_core.vectorstores import InMemoryVectorStore

vector_store = InMemoryVectorStore(embeddings)

我们现在可以构建一个检索工具，该工具可以在数据库中搜索相关的专有名词

from langchain.agents.agent_toolkits import create_retriever_tool

_ = vector_store.add_texts(artists + albums)
retriever = vector_store.as_retriever(search_kwargs={"k": 5})
description = (
    "Use to look up values to filter on. Input is an approximate spelling "
    "of the proper noun, output is valid proper nouns. Use the noun most "
    "similar to the search."
)
retriever_tool = create_retriever_tool(
    retriever,
    name="search_proper_nouns",
    description=description,
)

API 参考：create_retriever_tool

让我们试用一下

print(retriever_tool.invoke("Alice Chains"))

Alice In Chains

Alanis Morissette

Pearl Jam

Pearl Jam

Audioslave

这样，如果代理确定它需要根据“Alice Chains”之类的艺术家编写过滤器，它可以首先使用检索器工具来观察列的相关值。

将它们放在一起

# Add to system message
suffix = (
    "If you need to filter on a proper noun like a Name, you must ALWAYS first look up "
    "the filter value using the 'search_proper_nouns' tool! Do not try to "
    "guess at the proper name - use this function to find similar ones."
)

system = f"{system_message}\n\n{suffix}"

tools.append(retriever_tool)

agent = create_react_agent(llm, tools, prompt=system)

question = "How many albums does alis in chain have?"

for step in agent.stream(
    {"messages": [{"role": "user", "content": question}]},
    stream_mode="values",
):
    step["messages"][-1].pretty_print()

================================[1m Human Message [0m=================================

How many albums does alis in chain have?
==================================[1m Ai Message [0m==================================
Tool Calls:
  search_proper_nouns (call_8ryjsRPLAr79mM3Qvnq6gTOH)
 Call ID: call_8ryjsRPLAr79mM3Qvnq6gTOH
  Args:
    query: alis in chain
=================================[1m Tool Message [0m=================================
Name: search_proper_nouns

Alice In Chains

Aisha Duo

Xis

Da Lama Ao Caos

A-Sides
==================================[1m Ai Message [0m==================================
Tool Calls:
  sql_db_list_tables (call_NJjtCpU89MBMplssjn1z0xzq)
 Call ID: call_NJjtCpU89MBMplssjn1z0xzq
  Args:
  search_proper_nouns (call_1BfrueC9koSIyi4OfMu2Ao8q)
 Call ID: call_1BfrueC9koSIyi4OfMu2Ao8q
  Args:
    query: Alice In Chains
=================================[1m Tool Message [0m=================================
Name: search_proper_nouns

Alice In Chains

Pearl Jam

Pearl Jam

Foo Fighters

Soundgarden
==================================[1m Ai Message [0m==================================
Tool Calls:
  sql_db_schema (call_Kn09w9jd9swcNzIZ1b5MlKID)
 Call ID: call_Kn09w9jd9swcNzIZ1b5MlKID
  Args:
    table_names: Album, Artist
=================================[1m Tool Message [0m=================================
Name: sql_db_schema


CREATE TABLE "Album" (
	"AlbumId" INTEGER NOT NULL, 
	"Title" NVARCHAR(160) NOT NULL, 
	"ArtistId" INTEGER NOT NULL, 
	PRIMARY KEY ("AlbumId"), 
	FOREIGN KEY("ArtistId") REFERENCES "Artist" ("ArtistId")
)

/*
3 rows from Album table:
AlbumId	Title	ArtistId
1	For Those About To Rock We Salute You	1
2	Balls to the Wall	2
3	Restless and Wild	2
*/


CREATE TABLE "Artist" (
	"ArtistId" INTEGER NOT NULL, 
	"Name" NVARCHAR(120), 
	PRIMARY KEY ("ArtistId")
)

/*
3 rows from Artist table:
ArtistId	Name
1	AC/DC
2	Accept
3	Aerosmith
*/
==================================[1m Ai Message [0m==================================
Tool Calls:
  sql_db_query (call_WkHRiPcBoGN9bc58MIupRHKP)
 Call ID: call_WkHRiPcBoGN9bc58MIupRHKP
  Args:
    query: SELECT COUNT(*) FROM Album WHERE ArtistId = (SELECT ArtistId FROM Artist WHERE Name = 'Alice In Chains')
=================================[1m Tool Message [0m=================================
Name: sql_db_query

[(1,)]
==================================[1m Ai Message [0m==================================

Alice In Chains has released 1 album in the database.

正如我们所看到的，在流式传输的步骤和 LangSmith 跟踪中，代理都使用了 search_proper_nouns 工具，以检查如何正确查询数据库以查找此特定艺术家。

⚠️ 安全注意事项 ⚠️​

架构​

设置​

示例数据​

链​

应用程序状态​

将问题转换为 SQL 查询​

执行查询​

生成答案​

使用 LangGraph 编排​

人工参与环节​

后续步骤​

代理​

系统提示​

初始化代理​

处理高基数列​

此页面是否对您有帮助？