给大模型装上手，通过 Function Calling 查询数据库：实现高效数据处理的新方式

# 导入必要的库import openaiimport osfrom math import *from icecream import icimport jsonfrom math import *import requestsimport loggingimport sqlite3
# 设置日志记录配置logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# 加载 .env 文件from dotenv import load_dotenv, find_dotenv_ = load_dotenv(find_dotenv())
# 从环境变量中获得你的 OpenAI Keyopenai.api_key = os.getenv('OPENAI_API_KEY')openai.api_base = os.getenv('OPENAI_API_URL')model = os.getenv('MODEL')amap_key = os.getenv('GAODE_MAP_API_KEY')
# 定义数据库表结构的字符串database_schema_string = """CREATE TABLE customers (id INT PRIMARY KEY NOT NULL, -- 主键，不允许为空customer_name VARCHAR(255) NOT NULL, -- 客户名，不允许为空email VARCHAR(255) UNIQUE, -- 邮箱，唯一register_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP -- 注册时间，默认为当前时间);CREATE TABLE products (id INT PRIMARY KEY NOT NULL, -- 主键，不允许为空product_name VARCHAR(255) NOT NULL, -- 产品名称，不允许为空price DECIMAL(10,2) NOT NULL -- 价格，不允许为空);CREATE TABLE orders (id INT PRIMARY KEY NOT NULL, -- 主键，不允许为空customer_id INT NOT NULL, -- 客户ID，不允许为空product_id INT NOT NULL, -- 产品ID，不允许为空price DECIMAL(10,2) NOT NULL, -- 价格，不允许为空status INT NOT NULL, -- 订单状态，整数类型，不允许为空。0代表待支付，1代表已支付，2代表已退款create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, -- 创建时间，默认为当前时间pay_time TIMESTAMP -- 支付时间，可以为空);"""
def get_sql_completion(messages, model=model):"""使用 OpenAI 的 Chat API 来生成完成的 SQL 查询语句。参数:- messages: 一个包含用户消息和系统消息的列表，用于上下文对话。- model: 使用的 OpenAI 模型的名称，默认为环境变量中定义的模型。返回:- 生成的 SQL 查询语句。"""response = openai.ChatCompletion.create(model=model,messages=messages,temperature=0,# 模型输出的随机性，0 表示随机性最小tools=[{# 摘自 OpenAI 官方示例 https://github.com/openai/openai-cookbook/blob/main/examples/How_to_call_functions_with_chat_models.ipynb"type": "function","function": {"name": "ask_database","description": "Use this function to answer user questions about business. \Output should be a fully formed SQL query.",#使用此功能可以回答用户有关业务的问题\输出应该是完全格式的SQL查询。"parameters": {"type": "object","properties": {"query": {"type": "string",#SQL查询提取信息以回答用户的问题。SQL应该使用以下数据库架构编写：\#查询应以纯文本形式返回，而不是以JSON形式返回。查询应该只包含SQLite支持的语法。"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.The query should only contain grammars supported by SQLite.""",}},"required": ["query"],}}}],)return response.choices[0].message
# 创建内存中的SQLite数据库连接conn = sqlite3.connect(':memory:')cursor = conn.cursor()
# 执行数据库表创建语句for statement in database_schema_string.split(';'):cursor.execute(statement)
# 插入模拟数据mock_data = [(1, 1001, 'TSHIRT_1', 50.00, 0, '2023-10-12 10:00:00', None),(2, 1001, 'TSHIRT_2', 75.50, 1, '2023-10-16 11:00:00', '2023-08-16 12:00:00'),(3, 1002, 'SHOES_X2', 25.25, 2, '2023-10-17 12:30:00', '2023-08-17 13:00:00'),(4, 1003, 'HAT_Z112', 60.75, 1, '2023-10-20 14:00:00', '2023-08-20 15:00:00'),(5, 1002, 'WATCH_X001', 90.00, 0, '2023-10-28 16:00:00', None)]
for record in mock_data:cursor.execute('''\INSERT INTO orders (id, customer_id, product_id, price, status, create_time, pay_time)VALUES (?, ?, ?, ?, ?, ?, ?)''', record)
# 提交事务conn.commit()
def ask_database(query):"""执行 SQL 查询并返回结果。参数:- query: 要执行的 SQL 查询语句。返回:- 查询结果。"""cursor.execute(query)records = cursor.fetchall()return records
def test_promopt():"""测试通过自然语言询问来生成并执行 SQL 查询的功能。"""prompt = "统计每月每件商品的销售额"# prompt = "这星期消费最高的用户是谁？他买了哪些商品？ 每件商品买了几件？花费多少？"messages = [{"role": "system", "content": "基于 order 表回答用户问题"},{"role": "user", "content": prompt}]response = get_sql_completion(messages)print(response)conn.close()
def test_promopt_result():"""测试自然语言询问的更复杂场景，包括对数据库操作的调用和日志记录。"""prompt = "统计每月每件商品的销售额"# prompt = "这星期消费最高的用户是谁？他买了哪些商品？ 每件商品买了几件？花费多少？"messages = [{"role": "system", "content": "基于 order 表回答用户问题"},{"role": "user", "content": prompt}]response = get_sql_completion(messages)if response.content is None:response.content = "null"messages.append(response)logging.info("====Function Calling====")logging.info(response)
if response.tool_calls is not None:tool_call = response.tool_calls[0]if tool_call.function.name == "ask_database":arguments = tool_call.function.argumentsargs = json.loads(arguments)logging.info("====SQL====")logging.info(args["query"])result = ask_database(args["query"])logging.info("====DB Records====")logging.info(result)
messages.append({"tool_call_id": tool_call.id,"role": "tool","name": "ask_database","content": str(result)})response = get_sql_completion(messages)logging.info("====最终回复====")logging.info(response.content)conn.close()

if __name__ == '__main__': test_promopt_result()

2024-05-31 23:12:38,331 - INFO - ====Function Calling====2024-05-31 23:12:38,332 - INFO - {"role": "assistant","content": "null","tool_calls": [{"id": "call_zhket5tQwLB9jGukvB2GpfMX","type": "function","function": {"name": "ask_database","arguments": "{\"query\":\"SELECT strftime('%Y-%m', create_time) AS month, product_id, SUM(price) AS total_sales FROM orders WHERE status = 1 GROUP BY month, product_id;\"}"}}]}2024-05-31 23:12:38,334 - INFO - ====SQL====2024-05-31 23:12:38,334 - INFO - SELECT strftime('%Y-%m', create_time) AS month, product_id, SUM(price) AS total_sales FROM orders WHERE status = 1 GROUP BY month, product_id;2024-05-31 23:12:38,335 - INFO - ====DB Records====2024-05-31 23:12:38,335 - INFO - [('2023-10', 'HAT_Z112', 60.75), ('2023-10', 'TSHIRT_2', 75.5)]2024-05-31 23:12:41,284 - INFO - ====最终回复====2024-05-31 23:12:41,285 - INFO - 每月每件商品的销售额统计如下：
- 2023年10月：- 商品ID：HAT_Z112，销售额：60.75- 商品ID：TSHIRT_2，销售额：75.5