使用Azure OpenAI与Langchain生成SQL时Schema名称缺失的问题

我正在使用Azure OpenAI和Langchain实现Text to SQL功能，目标是从自然文本提取SQL查询。该功能基于CSV文件中的表详情（包含表名、表描述、Schema名字段）生成SQL，要求生成的SQL必须同时包含Schema名和表名，但目前生成的SQL未包含Schema名，请求排查解决。

代码

from operator import itemgetter
from langchain.chains.openai_tools import create_extraction_chain_pydantic
from langchain_core.pydantic_v1 import BaseModel, Field
from typing import List
import pandas as pd

def get_table_details():
    # Read the CSV file into a DataFrame
    table_description = pd.read_csv("database_table_descriptions.csv")
    table_docs = []

    # Iterate over the DataFrame rows to create Document objects
    table_details = ""
    for index, row in table_description.iterrows():
        table_details = table_details + "Table Name:" + row['table'] + "\n" + "Table Description:" + row['description'] +  "\n" + "Schema Name:" + row['schema'] +"\n\n"

    return table_details


class Table(BaseModel):
    """Table in SQL database."""

    name: str = Field(description="Name of table in SQL database.")
    schema_name: str = Field(description="schema of table in SQL database.")

table_details = get_table_details()
print(table_details)

table_details_prompt = f"""Return the names of ALL the SQL tables and names of ALL the SQL schema that MIGHT be relevant to the user question. \
The tables are:

{table_details}

Remember to include ALL POTENTIALLY RELEVANT tables, even if you're not sure that they're needed."""

table_chain = create_extraction_chain_pydantic(Table, llm, system_message=table_details_prompt)
def get_tables(tables: List[Table]) -> List[str]:
    tables  = [f"{table.schema_name}.{table.name}" for table in tables]
    return tables

select_table = {"input": itemgetter("question")} | create_extraction_chain_pydantic(Table, llm, system_message=table_details_prompt) | get_tables
chain = (
RunnablePassthrough.assign(schema_and_table_names_to_use=select_table) |
RunnablePassthrough.assign(query=generate_query).assign(
    result=itemgetter("query") | execute_query
)
)
res = chain.invoke({"question": "what is the ftpjndi for tradingpartner Rockwood for table qtc_asn"})

当前输出

{'question': 'what is the ftpjndi for tradingpartner Rockwood for table qtc_asn',
 'schema_and_table_names_to_use': ['mule.qtc_asn'],
 'query': 'SELECT "ftpjndi" FROM qtc_asn WHERE tradingpartner = \'RockwoodSyspro\' LIMIT 1;',
 'result': ''}

期望输出

{'question': 'what is the ftpjndi for tradingpartner Rockwood for table qtc_asn',
 'schema_and_table_names_to_use': ['mule.qtc_asn'],
 'query': 'SELECT "ftpjndi" FROM mule.qtc_asn WHERE tradingpartner = \'RockwoodSyspro\' LIMIT 1;',
 'result': 'abc/xyz'}

解决方案 #

1. 强化generate_query的提示词约束

问题核心在于生成SQL的环节未强制使用带Schema的完整表名。修改generate_query的提示模板，明确要求必须使用schema_and_table_names_to_use中的完整表名格式（schema.table）：

def generate_query(inputs):
    prompt = f"""根据用户问题生成符合要求的SQL查询：
    用户问题：{inputs['question']}
    必须使用以下指定的完整表名（包含Schema）：{inputs['schema_and_table_names_to_use']}
    注意：SQL中的表名必须严格采用`schema_name.table_name`格式，禁止仅使用表名。"""
    return llm.invoke(prompt).content

2. 确保generate_query正确接收参数

确认链的构建中，generate_query能正确访问到schema_and_table_names_to_use变量。当前链的结构已经传递了该参数，但需确保generate_query函数内部确实使用了这个参数来约束表名。

3. 优化表提取提示词（可选）

虽然当前表提取环节已正确返回带Schema的表名，但可以进一步强化提示词，确保返回的表信息明确关联Schema：

table_details_prompt = f"""返回所有与用户问题相关的SQL表名和对应的Schema名，必须以`schema_name.table_name`的格式返回所有可能相关的表，即使不确定是否需要。
可用表信息：
{table_details}"""

4. 验证数据库执行环境

如果数据库连接默认Schema为mule，LLM可能会省略Schema名，但根据需求必须强制添加，因此仍需通过提示词严格约束，避免依赖数据库默认配置。

内容的提问来源于stack exchange，提问作者Ganesh