QingLong/AI/Text2Sql/YunXiao.py

import os

from pyecharts import options as opts
from pyecharts.charts import Bar

from Text2Sql.Util.PostgreSQLUtil import PostgreSQLUtil
from Text2Sql.Util.SaveToExcel import save_to_excel
from Text2Sql.Util.Text2SqlUtil import *


def generate_bar_chart_from_result(result, x_column: str, y_columns: list, output_file: str = "bar_chart.html"):
    """
    根据结果集生成柱状图并保存为 HTML 文件。

    参数:
        result: 结果集，通常是一个列表（字典格式）。
        x_column (str): 作为 X 轴的列名。
        y_columns (list): 作为 Y 轴的列名列表。
        output_file (str): 输出 HTML 文件名（默认 "bar_chart.html"）。
    """
    try:
        # 如果结果集为空，直接返回
        if not result:
            raise ValueError("结果集为空，无法生成图表！")

        # 检查列名是否存在
        if x_column not in result[0]:
            raise ValueError(f"X 轴列名 '{x_column}' 不存在！")
        for col in y_columns:
            if col not in result[0]:
                raise ValueError(f"Y 轴列名 '{col}' 不存在！")

        # 提取 X 轴和 Y 轴数据
        x_axis_data = [row[x_column] for row in result]
        series_data = [
            (col, [row[col] for row in result])
            for col in y_columns
        ]

        print("X 轴数据:", x_axis_data)
        print("Y 轴数据:", series_data)

        # 创建柱状图
        bar = Bar()
        bar.add_xaxis(x_axis_data)
        for name, data in series_data:
            bar.add_yaxis(name, data)

        # 设置全局配置
        bar.set_global_opts(
            title_opts=opts.TitleOpts(title="动态柱状图"),
            xaxis_opts=opts.AxisOpts(name=x_column),
            yaxis_opts=opts.AxisOpts(name="值"),
            legend_opts=opts.LegendOpts(is_show=True)
        )

        # 确保目标目录存在
        if not os.path.exists(os.path.dirname(output_file)):
            os.makedirs(os.path.dirname(output_file))

        # 保存为 HTML 文件
        bar.render(output_file)
        print(f"柱状图已保存为 {output_file}")

        # 检查文件是否存在
        if os.path.exists(output_file):
            print("文件生成成功！")
        else:
            print("文件生成失败！")
    except Exception as e:
        print(f"生成柱状图时发生错误: {e}")


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

    # 开始训练
    print("开始训练...")
    # 打开CreateTable.sql文件内容
    with open("Sql/AreaSchoolLesson.sql", "r", encoding="utf-8") as file:
        ddl = file.read()
    # 训练数据
    vn.train(
        ddl=ddl
    )

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

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

    # 指定学段
    question = '''
                查询:
                1、发布时间是2024年度
                2、每个学段，每个科目，上传课程数量，按由多到少排序
                3、字段名: 学段,科目,排名,课程数量
                '''
    common_prompt = '''
    返回的信息要求：
    1、行政区划为NULL 或者是空字符的不参加统计
    2、目标数据库是Postgresql 16,注意字段名称不要有二义性问题
    3、使用stage_name描述学段,一定不要使用stage_id
    4、使用subject_name描述科目,一定不要使用subject_id
    '''
    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)

        # 生成柱状图
        # generate_bar_chart_from_result(
        #     _data,
        #     x_column="科目",
        #     y_columns=["课程数量"],
        #     output_file="d:/lesson_bar_chart.html"
        # )
        filename = "d:/导出信息.xlsx"
        save_to_excel(_data, filename)
        os.startfile(filename)