使用LLM实现自然语言的SQL查询代理

def connect_db() -> Connection:
    """创建sqlite数据库连接"""

    # 找到本地的chinook.db文件,创建数据库连接
    conn = sqlite3.connect("./chinook.db")
    print("connect to sqlite success!")
    return conn


def get_table_names(conn: Connection) -> List[str]:
    """返回一个包含所有表名的列表"""

    # 执行SQL查询,获取数据库中所有表的名字
    tables = conn.execute("SELECT name FROM sqlite_master WHERE type='table';")
    # 遍历查询table,返回table_name列表
    return [table[0] for table in tables]


def get_column_names(conn: Connection, table_name: str):
    """返回一个给定表的所有列名的列表"""

    # 执行SQL查询，获取表的所有列的信息
    columns = conn.execute(f"PRAGMA table_info('{table_name}');").fetchall()
    # 遍历columns,返回列名列表
    return [col[1] for col in columns]


def get_database_schema(conn: Connection) -> str:
    """获取数据库的Schema信息,包括表名及其包含的列表"""
    
    table_dicts = []  # 创建一个空的字典列表
    # 遍历数据库中的所有表
    for table_name in get_table_names(conn):
        columns_names = get_column_names(conn, table_name)  # 获取当前表的所有列名
        # 将表名和列名信息作为一个字典添加到列表中
        table_dicts.append({"table_name": table_name, "column_names": columns_names})

    # 格式化字典,转换成字符串返回
    database_schema_string = "\n".join(
        [
            f"Table: {table['table_name']}\nColumns: {', '.join(table['column_names'])}"
            for table in table_dicts
        ]
    )
    return database_schema_string
  
  def exec_sql(conn: Connection, sql: str) -> str:
    """执行SQL,获取结果"""
    
    try:
        # 执行SQL,并将结果转换为字符串
        results = str(conn.execute(sql).fetchall())
    except Exception as e:
        # 执行失败,捕获异常并返回错误信息
        results = f"query failed with error: {e}"

    # 返回查询结果
    return results

Table: albums
Columns: AlbumId, Title, ArtistId
Table: sqlite_sequence
Columns: name, seq
Table: artists
Columns: ArtistId, Name
Table: customers
Columns: CustomerId, FirstName, LastName, Company, Address, City, State, Country, PostalCode, Phone, Fax, Email, SupportRepId
Table: employees
Columns: EmployeeId, LastName, FirstName, Title, ReportsTo, BirthDate, HireDate, Address, City, State, Country, PostalCode, Phone, Fax, Email
Table: genres
Columns: GenreId, Name
Table: invoices
Columns: InvoiceId, CustomerId, InvoiceDate, BillingAddress, BillingCity, BillingState, BillingCountry, BillingPostalCode, Total
Table: invoice_items
Columns: InvoiceLineId, InvoiceId, TrackId, UnitPrice, Quantity
Table: media_types
Columns: MediaTypeId, Name
Table: playlists
Columns: PlaylistId, Name
Table: playlist_track
Columns: PlaylistId, TrackId
Table: tracks
Columns: TrackId, Name, AlbumId, MediaTypeId, GenreId, Composer, Milliseconds, Bytes, UnitPrice
Table: sqlite_stat1
Columns: tbl, idx, stat

def generate_sql(db_schema: str, query: str) -> str:
    """
    调用LLM,利用工具调用能力,生成SQL语句
    :param db_schema: 数据库表结构信息
    :param query: 用户的原始提问
    :return: 生成的结构化SQL
    """

    # 加载环境变量
    dotenv.load_dotenv()

    # 创建智谱AI客户端
    client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))

    # 定义工具的详细描述,便于LLM理解用户的需求
    tool_desc = f"""根据用户提问，生成的SQL语句，用于回答用户的问题。
                生成的SQL语句基于如下的数据库表结构定义：
                {db_schema}
                最终的SQL语句以纯文本的格式输出，不要使用json或者其它的结构化格式。
                """

    # 定义工具
    tools = [
        {
            "type": "function",  # 工具类型为function函数调用
            "function": {  # 函数定义
                "name": "generate_sql",  # 函数名称
                "description": "该函数用于回答用户提出的关于音乐的相关问题。 "
                               "生成的结果是结构化的标准SQL语句。",
                # 函数描述
                "parameters": {  # 函数参数定义
                    "type": "object",
                    "properties": {
                        "sql": {  # 参数名称
                            "type": "string",  # 参数类型
                            "description": tool_desc,  # 参数描述
                        },
                    },
                    "required": ["sql"],  # 必需的参数
                },
            }
        }
    ]

    # 创建消息列表
    messages = [
        {"role": "system",
         "content": "请根据用户的提问，基于Chinook Music数据库的信息，生成SQL语句来回答用户的问题。"},
        {"role": "user", "content": f"{query}"},
    ]

    # 执行工具调用,获取结果
    completion = client.chat.completions.create(
        model="glm-4-flash",
        messages=messages,
        tools=tools,
        tool_choice="auto"  # 工具选择模式为auto,表示由LLM自行推理,觉得是生成普通消息还是进行工具调用
    )

    # 将工具调用结果解析成sql字符串,并返回
    return json.loads(completion.choices[0].message.tool_calls[0].function.arguments).get("sql")

if __name__ == '__main__':
    # 创建sqlite数据库连接
    conn = connect_db()

    # 获取数据库的schema信息
    db_schema = get_database_schema(conn)

    # 在控制台循环获取用户输入
    while True:
        query = input("用户提问: ")
        if query == "bye":
            break

        # 将用户提问翻译成SQL
        sql = generate_sql(db_schema, query)
        print("--------------------------------------------------")
        print(f"生成的SQL语句: \n{sql}")

        # 执行SQL,获取结果
        answer = exec_sql(conn, sql)
        print("--------------------------------------------------")
        print(f"执行结果: {answer}")