如何使用chatGPT查询数据库中数据并进行数据分析

612 阅读1分钟
import mysql.connector
import openai
import json
from datetime import datetime
from quart import Quart, jsonify, send_file
from flask import Flask, jsonify, request
# app = Quart(__name__)
app = Flask(__name__)

class DateTimeEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.isoformat()
        return super().default(obj)

def query_database(sql):
    try:
        # 连接数据库
        host = "your host"
        user = "your username"
        password = "your password"
        database = "your database"
        connection = mysql.connector.connect(host=host, user=user, password=password, database=database)
        
        if connection.is_connected():
            print("成功连接到数据库!")
            # 创建游标对象
            cursor = connection.cursor()
            
            # 执行查询语句
            print(sql)
            cursor.execute(sql)
            
            # 获取查询结果
            result = cursor.fetchall()
            # 打印数据
            # for row in result:
            #     print(row)
            
            # 关闭游标和连接
            cursor.close()
            connection.close()
            json_data = json.dumps(result, cls=DateTimeEncoder)
            print(json_data)
            return json_data
            
    except mysql.connector.Error as e:
        print("查询数据库出错:", e)
        


# Step 1, send model the user query and what functions it has access to
def run_conversation(text):
    #  gpt 的key
    openai.api_key = ""  # Replace this with your OpenAI API key. This key can be found at
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo-0613",
        messages=[{"role": "user", "content": text}],
        functions=[
            {
                "name": "query_database",
                "description": text,
                "parameters": {
                    "type": "object",
                    "properties": {
                        "sql": {
                            "type": "string",
                            "description": "this is sql",
                        },
                        # "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
                    },
                    "required": ["sql"],
                },
            }
        ],
        function_call="auto",
    )
    print(response)
    message = response["choices"][0]["message"]
    print(message)
    # Step 2, check if the model wants to call a function
    if message.get("function_call"):
        function_name = message["function_call"]["name"]
        json_data = json.loads(message["function_call"]["arguments"])
        # Step 3, call the function
        # Note: the JSON response from the model may not be valid JSON
        function_response = query_database(
            sql=json_data["sql"],
            # unit=message.get("unit"),
        )

        # Step 4, send model the info on the function call and function response
        second_response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo-0613",
            messages=[
                {"role": "user", "content": text},
                message,
                {
                    "role": "function",
                    "name": function_name,
                    "content": function_response,
                },
            ],
        )
        return second_response
# text = "查询 basic_message 表中title字段是站内信测试的前2条数据?"
# print(run_conversation(text))

#  text = "查询 basic_message 表中所有title字段是站内信测试的数据?"
@app.route('/search', methods=['GET'])
def search():
    text = request.args.get('text')  # 获取查询文本
    response = run_conversation(text)  # 运行查询函数
    response = jsonify({'message': response})
    response.headers.add('Access-Control-Allow-Origin', '*')  # 设置跨域请求头
    return response

def main():
    app.run(debug=True, host="0.0.0.0", port=5002)


if __name__ == "__main__":
    main()

调用示例:

浏览器中直接调用/search 传参数:查询 your tablename 表中所有title字段是站内信测试的数据? image.png