定义一个执行SQL查询的函数

首先，让我们定义一些有用的实用函数来从SQLite数据库中提取数据。

step1 下载样本数据

www.sqlitetutorial.net/sqlite-samp…

step2 连接测试

import sqlite3

conn = sqlite3.connect("data/chinook.db")
print("Opened database successfully")

step3 首先定义三个函数

get_table_names、get_column_names和get_database_info 用于从数据库连接对象中

• 获取数据库的表名
• 表的列名
• 以及整体数据库的信息。

def get_column_names(conn, table_name):
    """返回一个给定表的所有列名的列表"""
    column_names = []  # 创建一个空的列名列表
    # 执行SQL查询，获取表的所有列的信息
    columns = conn.execute(f"PRAGMA table_info('{table_name}');").fetchall()
    # 遍历查询结果，并将每个列名添加到列表中
    for col in columns:
        column_names.append(col[1])
    return column_names  # 返回列名列表


def get_database_info(conn):
    """返回一个字典列表，每个字典包含一个表的名字和列信息"""
    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})
    return table_dicts  # 返回字典列表


def ask_database(conn, query):
    """使用 query 来查询 SQLite 数据库的函数。"""
    try:
        results = str(conn.execute(query).fetchall())  # 执行查询，并将结果转换为字符串
    except Exception as e:  # 如果查询失败，捕获异常并返回错误信息
        results = f"query failed with error: {e}"
    return results  # 返回查询结果

step4 定义一个功能列表

# 获取数据库信息，并存储为字典列表
database_schema_dict = get_database_info(conn)

# 将数据库信息转换为字符串格式，方便后续使用
database_schema_string = "\n".join(
    [
        f"Table: {table['table_name']}\nColumns: {', '.join(table['column_names'])}"
        for table in database_schema_dict
    ]
)


# 定义一个功能列表，其中包含一个功能字典，该字典定义了一个名为"ask_database"的功能，用于回答用户关于音乐的问题
functions = [
    {
        "name": "ask_database",
        "description": "Use this function to answer user questions about music. Output should be a fully formed SQL query.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": f"""
                            SQL query extracting info to answer the user's question.
                            SQL should be written using this database schema:
                            {database_schema_string}
                            The query should be returned in plain text, not in JSON.
                            """,
                }
            },
            "required": ["query"],
        },
    }
]

step5 定义查询函数(执行 SQL 查询)

定义两个函数ask_database和execute_function_call

• ask_database 用于实际执行 SQL 查询并返回结果
• execute_function_call 用于根据消息中的功能调用信息来执行相应的功能并获取结果

def ask_database(conn, query):
    """使用 query 来查询 SQLite 数据库的函数。"""
    try:
        results = str(conn.execute(query).fetchall())  # 执行查询，并将结果转换为字符串
    except Exception as e:  # 如果查询失败，捕获异常并返回错误信息
        results = f"query failed with error: {e}"
    return results  # 返回查询结果


def execute_function_call(message):
    """执行函数调用"""
    # 判断功能调用的名称是否为 "ask_database"
    if message["function_call"]["name"] == "ask_database":
        # 如果是，则获取功能调用的参数，这里是 SQL 查询
        query = json.loads(message["function_call"]["arguments"])["query"]
        # 使用 ask_database 函数执行查询，并获取结果
        results = ask_database(conn, query)
    else:
        # 如果功能调用的名称不是 "ask_database"，则返回错误信息
        results = f"Error: function {message['function_call']['name']} does not exist"
    return results  # 返回结果

step6 调用大模型

# 创建一个空的消息列表
messages = []

# 向消息列表中添加一个系统角色的消息，内容是 "Answer user questions by generating SQL queries against the Chinook Music Database."
messages.append({"role": "system", "content": "Answer user questions by generating SQL queries against the Chinook Music Database."})

# 向消息列表中添加一个用户角色的消息，内容是 "Hi, who are the top 5 artists by number of tracks?"
messages.append({"role": "user", "content": "Hi, who are the top 5 artists by number of tracks?"})

# 使用 chat_completion_request 函数获取聊天响应
chat_response = chat_completion_request(messages, functions)

# 从聊天响应中获取助手的消息
assistant_message = chat_response.json()["choices"][0]["message"]

# 将助手的消息添加到消息列表中
messages.append(assistant_message)

# 如果助手的消息中有功能调用
if assistant_message.get("function_call"):
    # 使用 execute_function_call 函数执行功能调用，并获取结果
    results = execute_function_call(assistant_message)
    # 将功能的结果作为一个功能角色的消息添加到消息列表中
    messages.append({"role": "function", "name": assistant_message["function_call"]["name"], "content": results})

# 打印对话
print(messages)

输出样例

assistant[function_call]: {'name': 'ask_database', 'arguments': '{"query":"SELECT ar.name AS Artist, COUNT(t.track_id) AS Number_of_Tracks\\nFROM artist ar\\nJOIN album al ON ar.artist_id = al.artist_id\\nJOIN track t ON al.album_id = t.album_id\\nGROUP BY ar.artist_id\\nORDER BY Number_of_Tracks DESC\\nLIMIT 5;"}'}

补充chat_completion_request定义

def chat_completion_request(messages, functions=None, function_call=None, model=GPT_MODEL):
    # 设定请求的header信息，包括 API_KEY
    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer " + os.getenv("OPENAI_API_KEY"),
    }

    # 设定请求的JSON数据，包括GPT 模型名和要进行补全的消息
    json_data = {"model": model, "messages": messages}

    # 如果传入了functions，将其加入到json_data中
    if functions is not None:
        json_data.update({"functions": functions})

    # 如果传入了function_call，将其加入到json_data中
    if function_call is not None:
        json_data.update({"function_call": function_call})

    # 尝试发送POST请求到OpenAI服务器的chat/completions接口
    try:
        response = requests.post(
            "https://api.openai.com/v1/chat/completions",
            headers=headers,
            json=json_data,
        )
        # 返回服务器的响应
        return response

    # 如果发送请求或处理响应时出现异常，打印异常信息并返回
    except Exception as e:
        print("Unable to generate ChatCompletion response")
        print(f"Exception: {e}")
        return e

总结

流程

• 定义好function函数
• 调用大模型 大模型返回结果 告诉你调用哪个函数 包括具体的参数值
• 调用函数 获取结果
• 将结果查询的结果放到message队列中 重新传给大模型