devbox
知新_RL
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

利用大模型实现数据库（sqlite）查询示例_大模型使用工具来查表查库的应用

作者：知新_RL | 2024-02-21 21:29:14

赞

踩

大模型使用工具来查表查库的应用
  1. import sqlite3
  2. import json
  3. import openai
  4. import requests
  5. import os
  6. from tenacity import retry, wait_random_exponential, stop_after_attempt
  7. from termcolor import colored
  8.  
  9. GPT_MODEL = "gpt-3.5-turbo-1106"
  10.  
  11. from openai import OpenAI
  12. client = OpenAI(
  13.     api_key = os.getenv("OPENAI_API_KEY"),
  14. )
  15.  
  16. # 1. 定义数据库信息
  17. conn = sqlite3.connect("data/chinook.db")
  18. print("Opened database successfully")
  19.  
  20. #--------------------------------------------------------
  21. # 定义一个函数chat_completion_request，主要用于发送 聊天补全 请求到OpenAI服务器
  22. def chat_completion_request(messages, tools=None):
  23.  
  24.      if tools is not None:
  25.         response = client.chat.completions.create(
  26.             model=GPT_MODEL,
  27.             messages=messages,
  28.             tools=tools,
  29.             tool_choice="auto",  # auto is default, but we'll be explicit
  30.         )
  31.      if tools is None:
  32.         response = client.chat.completions.create(
  33.                 model=GPT_MODEL,
  34.                 messages=messages
  35.             )
  36.  
  37.      return response
  38.  
  39.  
  40. # 定义一个函数pretty_print_conversation，用于打印消息对话内容
  41. def pretty_print_conversation(messages):
  42.  
  43.     # 为不同角色设置不同的颜色
  44.     role_to_color = {
  45.         "system": "red",
  46.         "user": "green",
  47.         "assistant": "blue",
  48.         "function": "magenta",
  49.     }
  50.  
  51.     # 遍历消息列表
  52.     for message in messages:
  53.         # 如果消息的角色是"system"，则用红色打印“content”
  54.         if message["role"] == "system":
  55.             print(colored(f"system: {message['content']}\n", role_to_color[message["role"]]))
  56.  
  57.         # 如果消息的角色是"user"，则用绿色打印“content”
  58.         elif message["role"] == "user":
  59.             print(colored(f"user: {message['content']}\n", role_to_color[message["role"]]))
  60.  
  61.         # 如果消息的角色是"assistant"，并且消息中包含"function_call"，则用蓝色打印"function_call"
  62.         elif message["role"] == "assistant" and message.get("function_call"):
  63.             print(colored(f"assistant[function_call]: {message['function_call']}\n", role_to_color[message["role"]]))
  64.  
  65.         # 如果消息的角色是"assistant"，但是消息中不包含"function_call"，则用蓝色打印“content”
  66.         elif message["role"] == "assistant" and not message.get("function_call"):
  67.             print(colored(f"assistant[content]: {message['content']}\n", role_to_color[message["role"]]))
  68.  
  69.         # 如果消息的角色是"function"，则用品红色打印“function”
  70.         elif message["role"] == "function":
  71.             print(colored(f"function ({message['name']}): {message['content']}\n", role_to_color[message["role"]]))
  72.  
  73.  
  74. def get_table_names(conn):
  75.     """返回一个包含所有表名的列表"""
  76.     table_names = []  # 创建一个空的表名列表
  77.     # 执行SQL查询，获取数据库中所有表的名字
  78.     tables = conn.execute("SELECT name FROM sqlite_master WHERE type='table';")
  79.     # 遍历查询结果，并将每个表名添加到列表中
  80.     for table in tables.fetchall():
  81.         table_names.append(table[0])
  82.     return table_names  # 返回表名列表
  83.  
  84.  
  85. def get_column_names(conn, table_name):
  86.     """返回一个给定表的所有列名的列表"""
  87.     column_names = []  # 创建一个空的列名列表
  88.     # 执行SQL查询，获取表的所有列的信息
  89.     columns = conn.execute(f"PRAGMA table_info('{table_name}');").fetchall()
  90.     # 遍历查询结果，并将每个列名添加到列表中
  91.     for col in columns:
  92.         column_names.append(col[1])
  93.     return column_names  # 返回列名列表
  94.  
  95.  
  96. def get_database_info(conn):
  97.     """返回一个字典列表，每个字典包含一个表的名字和列信息"""
  98.     table_dicts = []  # 创建一个空的字典列表
  99.     # 遍历数据库中的所有表
  100.     for table_name in get_table_names(conn):
  101.         columns_names = get_column_names(conn, table_name)  # 获取当前表的所有列名
  102.         # 将表名和列名信息作为一个字典添加到列表中
  103.         table_dicts.append({"table_name": table_name, "column_names": columns_names})
  104.     return table_dicts  # 返回字典列表
  105.  
  106.  
  107. #--------------------------------------------------------
  108. # 获取数据库信息，并存储为字典列表
  109. database_schema_dict = get_database_info(conn)
  110.  
  111. # 将数据库信息转换为字符串格式，方便后续使用
  112. database_schema_string = "\n".join(
  113.     [
  114.         f"Table: {table['table_name']}\nColumns: {', '.join(table['column_names'])}"
  115.         for table in database_schema_dict
  116.     ]
  117. )
  118.  
  119. # 2.定义 functions让 GPT 模型帮我们构造一个完整的 SQL 查询
  120. tools = [
  121.         {
  122.             "type": "function",
  123.             "function": {
  124.                 "name": "ask_database",
  125.                 "description": "Use this function to answer user questions about music. Output should be a fully formed SQL query.",
  126.                 "parameters": {
  127.                     "type": "object",
  128.                     "properties": {
  129.                         "query": {
  130.                             "type": "string",
  131.                             "description": f"""
  132.                             SQL query extracting info to answer the user's question.
  133.                             SQL should be written using this database schema:{database_schema_string}
  134.                             The query should be returned in plain text, not in JSON.
  135.                             """,
  136.                         } 
  137.                     },
  138.                     "required": ["query"],
  139.                 },
  140.             },
  141.         }
  142.     ]
  143.  
  144.  
  145. # 3. 执行SQL语句
  146.  
  147. def ask_database(conn, query):
  148.     """使用 query 来查询 SQLite 数据库的函数。"""
  149.     try:
  150.         results = str(conn.execute(query).fetchall())  # 执行查询，并将结果转换为字符串
  151.     except Exception as e:  # 如果查询失败，捕获异常并返回错误信息
  152.         results = f"query failed with error: {e}"
  153.     return results  # 返回查询结果
  154.  
  155.  
  156. def execute_function_call(message):
  157.     """执行函数调用"""
  158.     # 判断功能调用的名称是否为 "ask_database"
  159.     tool_calls = assistant_message.tool_calls
  160.     function_name = tool_calls[0].function.name
  161.     if function_name == "ask_database":
  162.         # 如果是，则获取功能调用的参数，这里是 SQL 查询
  163.         # query = json.loads(message["function_call"]["arguments"])["query"]
  164.         query = json.loads(tool_calls[0].function.arguments)["query"]
  165.         
  166.         # 使用 ask_database 函数执行查询，并获取结果
  167.         results = ask_database(conn, query)
  168.         print(query)
  169.         print(results)
  170.     else:
  171.         # 如果功能调用的名称不是 "ask_database"，则返回错误信息
  172.         results = f"Error: function {message['function_call']['name']} does not exist"
  173.     return results  # 返回结果
  174.  
  175.  
  176.  
  177. # 创建一个空的消息列表
  178. messages = []
  179.  
  180. # 向消息列表中添加一个系统角色的消息，内容是 "Answer user questions by generating SQL queries against the Chinook Music Database."
  181. messages.append({"role": "system",
  182.                  "content": "通过针对 Chinook 音乐数据库生成 SQL 查询来回答用户的问题"
  183.                  })
  184. # "content": "Answer user questions by generating SQL queries against the Chinook Music Database."
  185.  
  186. # 向消息列表中添加一个用户角色的消息，内容是 "Hi, who are the top 5 artists by number of tracks?"
  187. ## 根据曲目数量排名前五的艺术家是谁？
  188. messages.append({"role": "user", 
  189.                  "content": "嗨，歌曲数量排名前五的艺术家是谁？"})
  190. # "Hi, who are the top 5 artists by number of tracks?"
  191.  
  192. # 使用 chat_completion_request 函数获取聊天响应
  193. chat_response = chat_completion_request(messages, tools)
  194.  
  195. chat_response
  196.  
  197. # 从聊天响应中获取助手的消息
  198. assistant_message = chat_response.choices[0].message
  199. assistant_message
  200.  
  201. # 将助手的消息添加到消息列表中
  202. # messages.append(assistant_message)
  203. messages.append({
  204.     "role": assistant_message.role,  # 消息的角色是 
  205.     "content": ''  # 消息的内容 assistant_message.content
  206. })
  207.  
  208. # 如果助手的消息中有功能调用
  209. tool_calls = assistant_message.tool_calls
  210. if tool_calls:
  211.     tool_call = tool_calls[0]
  212.     # 使用 execute_function_call 函数执行功能调用，并获取结果
  213.     results = execute_function_call(assistant_message)
  214.  
  215.     # 将功能的结果作为一个功能角色的消息添加到消息列表中
  216.     messages.append({  
  217.                      "tool_call_id": tool_call.id,
  218.                      "role": "function", 
  219.                      "name": tool_call.function.name,  
  220.                      "content": results})
  221.  
  222.     #再次调用模型，将message对象给大模型
  223.     second_response = chat_completion_request(messages)
  224.     assistant_message = second_response.choices[0].message
  225.     messages.append({
  226.         "role": assistant_message.role,  # 消息的角色是 
  227.         "content": assistant_message.content  # 消息的内容 
  228.     })
  229.  
  230. # 使用 pretty_print_conversation 函数打印对话
  231. pretty_print_conversation(messages)
  232.  
  233. # ---------------------------------------------------------------------------------
  234. # 向消息列表中添加一个用户的问题，内容是 "What is the name of the album with the most tracks?"
  235. ## 哪张专辑的曲目数量最多？
  236. messages.append({"role": "user", 
  237.                  "content": "拥有最多曲目的专辑叫什么名字？"})  # What is the name of the album with the most tracks ?
  238.  
  239. # 使用 chat_completion_request 函数获取聊天响应
  240. chat_response = chat_completion_request(messages, tools)
  241.  
  242. # 从聊天响应中获取助手的消息
  243. assistant_message = chat_response.choices[0].message
  244. assistant_message
  245.  
  246.  
  247. # 将助手的消息添加到消息列表中
  248. # messages.append(assistant_message)
  249. messages.append({"role": assistant_message.role, 
  250.                  "content": ''})
  251.  
  252. # 如果助手的消息中有功能调用
  253. if assistant_message.tool_calls:
  254.     tool_calls = assistant_message.tool_calls
  255.     tool_call = tool_calls[0]
  256.     # 使用 execute_function_call 函数执行功能调用，并获取结果
  257.     results = execute_function_call(assistant_message)
  258.     # 将功能的结果作为一个功能角色的消息添加到消息列表中
  259.     messages.append({ "tool_call_id": tool_call.id,
  260.                      "role": "function", 
  261.                      "name": tool_call.function.name,  
  262.                      "content": results})
  263.     #再次调用模型，将message对象给大模型
  264.     second_response = chat_completion_request(messages)
  265.     assistant_message = second_response.choices[0].message
  266.     messages.append({
  267.         "role": assistant_message.role,  # 消息的角色是 
  268.         "content": assistant_message.content  # 消息的内容 
  269.     })
  270.  
  271. # 使用 pretty_print_conversation 函数打印对话
  272. pretty_print_conversation(messages)
  273.  
  274.  
  275.  
  276.  
  277.  
  278.  
  279.  
  280.  
  281.  
  282.  
最终messages打印的内容：

  1. [{'role': 'system', 'content': '通过针对 Chinook 音乐数据库生成 SQL 查询来回答用户的问题'},
  2.  {'role': 'user', 'content': '嗨，歌曲数量排名前五的艺术家是谁？'},
  3.  {'role': 'assistant', 'content': ''},
  4.  {'tool_call_id': 'call_D8CgPbKLVuYJrFbi6om7zsam',
  5.   'role': 'function',
  6.   'name': 'ask_database',
  7.   'content': "[('Iron Maiden', 213), ('U2', 135), ('Led Zeppelin', 114), ('Metallica', 112), ('Lost', 92)]"},
  8.  {'role': 'assistant', 'content': ''},
  9.  {'tool_call_id': 'call_D8CgPbKLVuYJrFbi6om7zsam',
  10.   'role': 'function',
  11.   'name': 'ask_database',
  12.   'content': "[('Iron Maiden', 213), ('U2', 135), ('Led Zeppelin', 114), ('Metallica', 112), ('Lost', 92)]"},
  13.  {'role': 'assistant',
  14.   'content': '歌曲数量排名前五的艺术家分别是 Iron Maiden（213首歌曲）、U2（135首歌曲）、Led Zeppelin（114首歌曲）、Metallica（112首歌曲）和Lost（92首歌曲）。'},
  15.  {'role': 'user', 'content': '拥有最多曲目的专辑叫什么名字？'},
  16.  {'role': 'assistant', 'content': ''},
  17.  {'tool_call_id': 'call_M72g22u0r0dEX5BX3MmZT2lK',
  18.   'role': 'function',
  19.   'name': 'ask_database',
  20.   'content': "[('Greatest Hits', 57)]"},
  21.  {'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/知新_RL/article/detail/126688?site
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号