利用大模型实现数据库（sqlite）查询示例_大模型查询数据库

作者：小桥流水78 | 2024-07-31 20:22:54

踩

大模型查询数据库


import sqlite3
import json
import openai
import requests
import os
from tenacity import retry, wait_random_exponential, stop_after_attempt
from termcolor import colored
 
GPT_MODEL = "gpt-3.5-turbo-1106"
 
from openai import OpenAI
client = OpenAI(
    api_key = os.getenv("OPENAI_API_KEY"),
)
 
# 1. 定义数据库信息
conn = sqlite3.connect("data/chinook.db")
print("Opened database successfully")
 
#--------------------------------------------------------
# 定义一个函数chat_completion_request，主要用于发送 聊天补全 请求到OpenAI服务器
def chat_completion_request(messages, tools=None):
 
     if tools is not None:
        response = client.chat.completions.create(
            model=GPT_MODEL,
            messages=messages,
            tools=tools,
            tool_choice="auto",  # auto is default, but we'll be explicit
        )
     if tools is None:
        response = client.chat.completions.create(
                model=GPT_MODEL,
                messages=messages
            )
 
     return response
 
 
# 定义一个函数pretty_print_conversation，用于打印消息对话内容
def pretty_print_conversation(messages):
 
    # 为不同角色设置不同的颜色
    role_to_color = {
        "system": "red",
        "user": "green",
        "assistant": "blue",
        "function": "magenta",
    }
 
    # 遍历消息列表
    for message in messages:
        # 如果消息的角色是"system"，则用红色打印“content”
        if message["role"] == "system":
            print(colored(f"system: {message['content']}\n", role_to_color[message["role"]]))
 
        # 如果消息的角色是"user"，则用绿色打印“content”
        elif message["role"] == "user":
            print(colored(f"user: {message['content']}\n", role_to_color[message["role"]]))
 
        # 如果消息的角色是"assistant"，并且消息中包含"function_call"，则用蓝色打印"function_call"
        elif message["role"] == "assistant" and message.get("function_call"):
            print(colored(f"assistant[function_call]: {message['function_call']}\n", role_to_color[message["role"]]))
 
        # 如果消息的角色是"assistant"，但是消息中不包含"function_call"，则用蓝色打印“content”
        elif message["role"] == "assistant" and not message.get("function_call"):
            print(colored(f"assistant[content]: {message['content']}\n", role_to_color[message["role"]]))
 
        # 如果消息的角色是"function"，则用品红色打印“function”
        elif message["role"] == "function":
            print(colored(f"function ({message['name']}): {message['content']}\n", role_to_color[message["role"]]))
 
 
def get_table_names(conn):
    """返回一个包含所有表名的列表"""
    table_names = []  # 创建一个空的表名列表
    # 执行SQL查询，获取数据库中所有表的名字
    tables = conn.execute("SELECT name FROM sqlite_master WHERE type='table';")
    # 遍历查询结果，并将每个表名添加到列表中
    for table in tables.fetchall():
        table_names.append(table[0])
    return table_names  # 返回表名列表
 
 
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  # 返回字典列表
 
 
#--------------------------------------------------------
# 获取数据库信息，并存储为字典列表
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
    ]
)
 
# 2.定义 functions让 GPT 模型帮我们构造一个完整的 SQL 查询
tools = [
        {
            "type": "function",
            "function": {
                "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"],
                },
            },
        }
    ]
 
 
# 3. 执行SQL语句
 
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"
    tool_calls = assistant_message.tool_calls
    function_name = tool_calls[0].function.name
    if function_name == "ask_database":
        # 如果是，则获取功能调用的参数，这里是 SQL 查询
        # query = json.loads(message["function_call"]["arguments"])["query"]
        query = json.loads(tool_calls[0].function.arguments)["query"]
        
        # 使用 ask_database 函数执行查询，并获取结果
        results = ask_database(conn, query)
        print(query)
        print(results)
    else:
        # 如果功能调用的名称不是 "ask_database"，则返回错误信息
        results = f"Error: function {message['function_call']['name']} does not exist"
    return results  # 返回结果
 
 
 
# 创建一个空的消息列表
messages = []
 
# 向消息列表中添加一个系统角色的消息，内容是 "Answer user questions by generating SQL queries against the Chinook Music Database."
messages.append({"role": "system",
                 "content": "通过针对 Chinook 音乐数据库生成 SQL 查询来回答用户的问题"
                 })
# "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, tools)
 
chat_response
 
# 从聊天响应中获取助手的消息
assistant_message = chat_response.choices[0].message
assistant_message
 
# 将助手的消息添加到消息列表中
# messages.append(assistant_message)
messages.append({
    "role": assistant_message.role,  # 消息的角色是 
    "content": ''  # 消息的内容 assistant_message.content
})
 
# 如果助手的消息中有功能调用
tool_calls = assistant_message.tool_calls
if tool_calls:
    tool_call = tool_calls[0]
    # 使用 execute_function_call 函数执行功能调用，并获取结果
    results = execute_function_call(assistant_message)
 
    # 将功能的结果作为一个功能角色的消息添加到消息列表中
    messages.append({  
                     "tool_call_id": tool_call.id,
                     "role": "function", 
                     "name": tool_call.function.name,  
                     "content": results})
 
    #再次调用模型，将message对象给大模型
    second_response = chat_completion_request(messages)
    assistant_message = second_response.choices[0].message
    messages.append({
        "role": assistant_message.role,  # 消息的角色是 
        "content": assistant_message.content  # 消息的内容 
    })
 
# 使用 pretty_print_conversation 函数打印对话
pretty_print_conversation(messages)
 
# ---------------------------------------------------------------------------------
# 向消息列表中添加一个用户的问题，内容是 "What is the name of the album with the most tracks?"
## 哪张专辑的曲目数量最多？
messages.append({"role": "user", 
                 "content": "拥有最多曲目的专辑叫什么名字？"})  # What is the name of the album with the most tracks ?
 
# 使用 chat_completion_request 函数获取聊天响应
chat_response = chat_completion_request(messages, tools)
 
# 从聊天响应中获取助手的消息
assistant_message = chat_response.choices[0].message
assistant_message
 
 
# 将助手的消息添加到消息列表中
# messages.append(assistant_message)
messages.append({"role": assistant_message.role, 
                 "content": ''})
 
# 如果助手的消息中有功能调用
if assistant_message.tool_calls:
    tool_calls = assistant_message.tool_calls
    tool_call = tool_calls[0]
    # 使用 execute_function_call 函数执行功能调用，并获取结果
    results = execute_function_call(assistant_message)
    # 将功能的结果作为一个功能角色的消息添加到消息列表中
    messages.append({ "tool_call_id": tool_call.id,
                     "role": "function", 
                     "name": tool_call.function.name,  
                     "content": results})
    #再次调用模型，将message对象给大模型
    second_response = chat_completion_request(messages)
    assistant_message = second_response.choices[0].message
    messages.append({
        "role": assistant_message.role,  # 消息的角色是 
        "content": assistant_message.content  # 消息的内容 
    })
 
# 使用 pretty_print_conversation 函数打印对话
pretty_print_conversation(messages)

最终messages打印的内容：


[{'role': 'system', 'content': '通过针对 Chinook 音乐数据库生成 SQL 查询来回答用户的问题'},
 {'role': 'user', 'content': '嗨，歌曲数量排名前五的艺术家是谁？'},
 {'role': 'assistant', 'content': ''},
 {'tool_call_id': 'call_D8CgPbKLVuYJrFbi6om7zsam',
  'role': 'function',
  'name': 'ask_database',
  'content': "[('Iron Maiden', 213), ('U2', 135), ('Led Zeppelin', 114), ('Metallica', 112), ('Lost', 92)]"},
 {'role': 'assistant', 'content': ''},
 {'tool_call_id': 'call_D8CgPbKLVuYJrFbi6om7zsam',
  'role': 'function',
  'name': 'ask_database',
  'content': "[('Iron Maiden', 213), ('U2', 135), ('Led Zeppelin', 114), ('Metallica', 112), ('Lost', 92)]"},
 {'role': 'assistant',
  'content': '歌曲数量排名前五的艺术家分别是 Iron Maiden（213首歌曲）、U2（135首歌曲）、Led Zeppelin（114首歌曲）、Metallica（112首歌曲）和Lost（92首歌曲）。'},
 {'role': 'user', 'content': '拥有最多曲目的专辑叫什么名字？'},
 {'role': 'assistant', 'content': ''},
 {'tool_call_id': 'call_M72g22u0r0dEX5BX3MmZT2lK',
  'role': 'function',
  'name': 'ask_database',
  'content': "[('Greatest Hits', 57)]"},
 {'role': 'assistant', 'content': '拥有最多曲目的专辑名字是《Greatest Hits》。'}]

可以看出，第一次和大模型交互时，大模型返回的role是assistant，content是空的，tool_calls是有内容的

执行了tool_call后，将获取内容放到message中，再将message提交给大模型进行第二次交互，role是assistant，content是回答的内容，tool_calls是空的了。

可见大模型重新整理问题和答案后，重新输出了内容，看起来舒服多了。

system: 通过针对 Chinook 音乐数据库生成 SQL 查询来回答用户的问题

user: 嗨，歌曲数量排名前五的艺术家是谁？

assistant[content]: 

function (ask_database): [('Iron Maiden', 213), ('U2', 135), ('Led Zeppelin', 114), ('Metallica', 112), ('Lost', 92)]

assistant[content]: 歌曲数量排名前五的艺术家分别是 Iron Maiden（213首歌曲）、U2（135首歌曲）、Led Zeppelin（114首歌曲）、Metallica（112首歌曲）和Lost（92首歌曲）。

user: 拥有最多曲目的专辑叫什么名字？

assistant[content]: 

function (ask_database): [('Greatest Hits', 57)]

assistant[content]: 拥有最多曲目的专辑名字是《Greatest Hits》。

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/小桥流水78/article/detail/910557