QingLong/AI/Text2Sql/Model/biModel.py

import asyncpg


# 删除数据
async def delete_question(db: asyncpg.Connection, question_id: str):
    delete_sql = """
        DELETE FROM t_bi_question WHERE id = $1
    """
    await db.execute(delete_sql, question_id)


# 插入数据
async def insert_question(db: asyncpg.Connection, question_id: str, question: str):
    insert_sql = """
        INSERT INTO t_bi_question (id, question, state_id, is_system, is_collect)
        VALUES ($1, $2, $3, $4, $5)
    """
    await db.execute(insert_sql, question_id, question, 0, 0, 0)


# 修改数据
async def update_question_by_id(db: asyncpg.Connection, question_id: str, **kwargs):
    update_fields = {k: v for k, v in kwargs.items() if v is not None}
    if not update_fields:
        return False

    set_clause = ", ".join([f"{field} = ${i + 1}" for i, field in enumerate(update_fields.keys())])
    sql = f"""
        UPDATE t_bi_question
        SET {set_clause}
        WHERE id = ${len(update_fields) + 1}
    """
    params = list(update_fields.values()) + [question_id]

    try:
        await db.execute(sql, *params)
        return True
    except Exception as e:
        print(f"更新失败: {e}")
        return False


# 根据问题 ID 查询 SQL
async def get_question_by_id(db: asyncpg.Connection, question_id: str):
    select_sql = """
        SELECT * FROM t_bi_question WHERE id = $1
    """
    _data = await db.fetch(select_sql, question_id)
    return _data


# 根据 SQL 查询数据
async def get_data_by_sql(db: asyncpg.Connection, sql: str):
    _data = await db.fetch(sql)
    return _data


# 保存系统推荐
async def set_system_recommend_questions(db: asyncpg.Connection, question_id: str, flag: str):
    sql = """
        UPDATE t_bi_question
        SET is_system = $1 WHERE id = $2
    """
    try:
        await db.execute(sql, int(flag), question_id)
        return True
    except Exception as e:
        print(f"更新失败: {e}")
        return False


# 设置用户收藏
async def set_user_collect_questions(db: asyncpg.Connection, question_id: str, flag: str):
    sql = """
        UPDATE t_bi_question
        SET is_collect = $1 WHERE id = $2
    """
    try:
        await db.execute(sql, int(flag), question_id)
        return True
    except Exception as e:
        print(f"更新失败: {e}")
        return False


# 查询系统推荐问题
async def get_system_recommend_questions(db: asyncpg.Connection, offset: int, limit: int):
    query = """
        SELECT *
        FROM t_bi_question where is_system=1  ORDER BY id DESC  LIMIT $1 OFFSET $2;
    """
    return await db.fetch(query, limit, offset)


async def get_system_recommend_questions_count(db: asyncpg.Connection):
    query = """
        SELECT COUNT(*)
        FROM t_bi_question where is_system=1;
    """
    return await db.fetchval(query)


async def get_user_publish_questions(db: asyncpg.Connection, type_id: int, offset: int, limit: int):
    # 基础查询
    query = """
        SELECT *
        FROM t_bi_question
    """

    # 根据 type_id 动态添加 WHERE 条件
    if type_id == 1:
        query += " WHERE is_collect = 1"

    # 添加排序和分页
    query += " ORDER BY id DESC LIMIT $1 OFFSET $2;"

    # 执行查询
    return await db.fetch(query, limit, offset)


async def get_user_publish_questions_count(db: asyncpg.Connection):
    query = """
        SELECT COUNT(*) FROM t_bi_question;
    """
    return await db.fetchval(query)


# 获取数据集的字段名称
async def get_column_names(db: asyncpg.Connection, sql: str):
    # 执行查询（添加 LIMIT 1）
    sql = sql.replace(";", "")
    sql = sql + ' limit 1'
    result = await db.fetchrow(sql)
    # 获取列名
    # 获取列名
    if result:
        column_names = list(result.keys())
        return column_names
    else:
        return []


from openai import OpenAI

from Config import MODEL_NAME, MODEL_API_KEY

# 初始化 OpenAI 客户端
client = OpenAI(
    api_key=MODEL_API_KEY,
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)


# 利用AI获取数据集的X轴和Y轴的列名
def generate_columns_with_ai(data):
    """
    利用大模型解析数据，生成 category_columns_str 和 value_column_str
    :param data: 数据集（列表字典格式，例如：[{"行政区划名": "二道区", "学校名称": "清华附中", "课程数量": 100}, ...]）
    :param api_key: OpenAI API 密钥
    :return: (category_columns_str, value_column_str)
    """
    # 获取所有字段名
    columns = list(data[0].keys()) if data else []

    # 构造提示词
    prompt = f"""
    给定以下字段名列表：{columns}，请分析并回答以下问题：
    1. 哪些字段适合作为分类字段（category_columns_str）？请用逗号分隔。
    2. 哪个字段适合作为数值字段（value_column_str）？
    3. 以JSON格式返回结果,但不要输出 json``` ，```

    返回格式：
    category_columns_str: <字段1,字段2,...>
    value_column_str: <字段>
    """

    # 调用大模型
    response = client.chat.completions.create(
        model=MODEL_NAME,
        messages=[
            {"role": "system",
             "content": "你是一个专业的语义分类助手。"},
            {"role": "user", "content": prompt}
        ],
        max_tokens=500
    )

    # 解析模型返回的结果
    result = response.choices[0].message.content
    category_columns_str = result.split("category_columns_str: ")[1].split("\n")[0].strip()
    value_column_str = result.split("value_column_str: ")[1].strip()

    return category_columns_str, value_column_str