QingLong/AI/Text2Sql/YunXiao.py

import json

from openai import OpenAI

from Text2Sql.Util.MarkdownToDocxUtil import markdown_to_docx
from Text2Sql.Util.PostgreSQLUtil import PostgreSQLUtil
from Text2Sql.Util.SaveToExcel import save_to_excel
from Text2Sql.Util.VannaUtil import *

from Util.EchartsUtil import *

'''
经验：
1、尽量使用宽表，少用关联，越少越好
2、应该有一些固定的组合用法预置出来，给出范例，让用户可以简单修改后就能使用
3、应该有类似于 保存为用例，查询历史等功能，让用户方便利旧。
'''

if __name__ == "__main__":
    vn = VannaUtil()

    # 开始训练
    print("开始训练...")
    # 打开AreaSchoolLesson.sql文件内容
    with open("Sql/AreaSchoolLessonDDL.sql", "r", encoding="utf-8") as file:
        ddl = file.read()
    # 训练数据
    vn.train(
        ddl=ddl
    )
    # 添加有关业务术语或定义的文档
    # vn.train(documentation="Sql/AreaSchoolLesson.md")

    # 使用 SQL 进行训练
    with open('Sql/AreaSchoolLessonGenerate.sql', 'r', encoding='utf-8') as file:
        sql_content = file.read()
    # 使用正则表达式提取注释和 SQL 语句
    sql_pattern = r'/\*(.*?)\*/(.*?);'
    sql_snippets = re.findall(sql_pattern, sql_content, re.DOTALL)

    # 打印提取的注释和 SQL 语句
    for i, (comment, sql) in enumerate(sql_snippets, 1):
        vn.train(sql=comment.strip() + '\n' + sql.strip() + '\n')

    # 自然语言提问
    # 整体情况
    # question = '''
    #     查询:
    #     1、发布时间是2024年度
    #     2、每个行政区每个学校都上传了多少课程数量
    #     3、格式: 行政区划名,学段,排名,学校名称,课程数量
    #     '''

    # 指定行政区域
    # question = '''
    #         查询:
    #         1、发布时间是2024年度
    #         2、二道区每个学校都上传了多少课程数量
    #         3、格式: 行政区划名,学段,排名,学校名称,发布年份,课程数量
    #         '''

    # 指定学段
    question = '''
                查询:
                1、发布时间是2024年度
                2、每个学段，每个科目，上传课程数量，按由多到少排序
                3、字段名: 学段,科目,排名,课程数量
                '''
    common_prompt = '''
    返回的信息要求：
    1、行政区划为NULL 或者是空字符的不参加统计
    2、目标数据库是Postgresql 16
    '''
    question = question + common_prompt
    # 开始查询
    print("开始查询...")
    # 获取完整 SQL
    sql = vn.generate_sql(question)
    print("生成的查询 SQL:\n", sql)

    # 执行SQL查询
    with PostgreSQLUtil() as db:
        _data = db.execute_query(sql)

        # 获取字段名和数据示例
        field_names = list(_data[0].keys()) if _data else []
        sample_data = _data[:3]  # 取前 3 行作为示例数据

        # 1、生成柱状图
        generate_bar_chart(
            _data=_data,
            title="学段+科目课程数量柱状图",
            x_columns=['学段', '科目'],  # 动态指定 X 轴列
            y_columns=['课程数量'],  # 动态指定 Y 轴列
            output_file="d:/lesson_bar_chart.html"
        )
        # 2、生成饼状图
        generate_pie_chart(
            _data=_data,
            title="学段+科目分布",
            category_columns=['学段', '科目'],  # 多列组合参数
            value_column='课程数量',
            output_file="d:/lesson_pie_chart.html"
        )

        # 3、生成excel
        filename = "d:/导出信息.xlsx"
        save_to_excel(_data, filename)

        # 4、生成word报告
        prompt = '''
        请根据以下 JSON 数据，整理出2000字左右的话描述当前数据情况。要求：
        1、以Markdown格式返回，我将直接通过markdown格式生成Word。
        2、标题统一为：长春云校数据分析报告
        3、内容中不要提到JSON数据，统一称：数据
        4、尽量以条目列出，这样更清晰
        5、数据：
        '''
        prompt = prompt + json.dumps(_data, ensure_ascii=False)

        # 初始化 OpenAI 客户端
        client = OpenAI(
            api_key=MODEL_API_KEY,
            base_url=MODEL_API_URL,
        )

        # 调用 OpenAI API 生成总结（流式输出）
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": "你是一个数据分析助手，擅长从 JSON 数据中提取关键信息并生成详细的总结。"},
                {"role": "user", "content": prompt}
            ],
            max_tokens=3000,  # 控制生成内容的长度
            temperature=0.7,  # 控制生成内容的创造性
            stream=True  # 启用流式输出
        )

        # 初始化变量用于存储流式输出的内容
        summary = ""

        # 处理流式输出
        for chunk in response:
            if chunk.choices[0].delta.content:  # 检查是否有内容
                chunk_content = chunk.choices[0].delta.content
                print(chunk_content, end="", flush=True)  # 实时打印到控制台
                summary += chunk_content  # 将内容拼接到 summary 中

        # 保存markdown
        with open("d:/report.md", "w", encoding="utf-8") as file:
            file.write(summary)

        # 最终 summary 为完整的 Markdown 内容
        print("\n\n流式输出完成，summary 已拼接为完整字符串。")
        # 生成 Word 文档
        markdown_to_docx(summary, output_file="d:/report.docx")