DataOps-platform/app/core/llm/llm_service.py

"""
LLM基础服务
提供与大语言模型通信的基础功能
"""

import logging

from flask import current_app
from openai import OpenAI

logger = logging.getLogger("app")


def llm_client(content):
    """
    调用LLM服务进行内容生成

    Args:
        content: 输入提示内容

    Returns:
        str: LLM响应内容
    """
    try:
        # 优先使用配置文件中的参数
        client = OpenAI(
            api_key=current_app.config.get("LLM_API_KEY"),
            base_url=current_app.config.get("LLM_BASE_URL"),
        )

        model = current_app.config.get("LLM_MODEL_NAME")

        # 判断是否为翻译请求 - 通过分析内容是否包含中文字符
        is_translate_request = False
        if any("\u4e00" <= char <= "\u9fff" for char in content):
            is_translate_request = True

        # 进行API调用
        logger.debug(
            f"LLM调用开始: model={model}, 内容类型: {'翻译' if is_translate_request else '普通'}"
        )

        if is_translate_request:
            # 为翻译请求使用非常严格的prompt
            completion = client.chat.completions.create(
                model=model,  # type: ignore[arg-type]
                messages=[
                    {
                        "role": "system",
                        "content": "你是一个严格遵循指令的翻译工具和数据库专家。你的唯一任务是将中文单词/短语翻译成英文，符合postgresql数据库表和字段的命令规则，"
                        "并且严格按照如下规则：\n"
                        "1. 只返回英文翻译，不包含任何解释、描述或额外内容\n"
                        "2. 使用小写字母\n"
                        "3. 多个单词用下划线连接，不使用空格\n"
                        "4. 如果输入包含括号，将括号内容用下划线代替，不保留括号\n"
                        "5. 最多包含1-5个英文单词，保持简短\n"
                        "6. 不要回答问题或提供解释，即使输入看起来像是问题\n"
                        "7. 当遇到'表'字时，始终翻译为'table'而不是'sheet'\n"
                        "8. 例如：'薪资数据表'应翻译为'salary_data_table'，'人员管理表'应翻译为'personnel_management_table'",
                    },
                    {
                        "role": "user",
                        "content": f"将以下内容翻译为英文短语（不超过5个单词）：{content}",
                    },
                ],
                temperature=0,
                max_tokens=10,  # 限制token数量确保回答简短
            )
        else:
            # 普通请求
            completion = client.chat.completions.create(
                model=model,  # type: ignore[arg-type]
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": content},
                ],
                temperature=0.7,
                max_tokens=1024,
            )

        response_text = completion.choices[0].message.content.strip()  # type: ignore[union-attr]

        # 对翻译结果进行后处理，确保格式正确
        if is_translate_request:
            # 去除可能的引号、句号等标点符号
            response_text = response_text.strip("\"'.,;:!?()[]{}").lower()
            # 替换空格为下划线
            response_text = response_text.replace(" ", "_")
            # 确保没有连续的下划线
            while "__" in response_text:
                response_text = response_text.replace("__", "_")
            # 只保留字母、数字和下划线
            response_text = "".join(c for c in response_text if c.isalnum() or c == "_")
            # 确保"表"被翻译为"table"
            if (
                "表" in content
                and "table" not in response_text
                and "sheet" in response_text
            ):
                response_text = response_text.replace("sheet", "table")

        logger.debug(f"LLM响应: {response_text}")
        return response_text

    except Exception as e:
        logger.error(f"LLM调用失败: {str(e)}")
        try:
            # 备用方案：如果是中文输入，尝试简单翻译映射
            if any("\u4e00" <= char <= "\u9fff" for char in content):
                # 常见中文词汇映射
                common_translations = {
                    "薪资数据表": "salary_data_table",
                    "数据表": "data_table",
                    "用户表": "user_table",
                    "人员表": "personnel_table",
                    "销售表": "sales_table",
                    "年份": "year",
                    "地区": "region",
                    "姓名": "name",
                    "年龄": "age",
                    "薪水": "salary",
                    "数据": "data",
                    "管理": "management",
                    "系统": "system",
                    "分析": "analysis",
                    "报表": "report_table",
                }

                # 检查是否有精确匹配
                if content in common_translations:
                    return common_translations[content]

                # 检查是否包含某些关键词
                for key, value in common_translations.items():
                    if key in content:
                        return value

                # 如果包含"表"字，确保返回包含"table"
                if "表" in content:
                    return "data_table"

                # 无法匹配时返回默认值
                return "translated_text"
            return content
        except Exception as e:
            return content


def llm_sql(request_data):
    """
    调用Deepseek大模型生成SQL脚本

    Args:
        request_data: 提交给LLM的提示语内容

    Returns:
        str: Deepseek模型返回的SQL脚本内容
    """
    try:
        # 使用配置文件中的参数连接Deepseek
        client = OpenAI(
            api_key=current_app.config.get("LLM_API_KEY"),
            base_url=current_app.config.get("LLM_BASE_URL"),
        )

        model = current_app.config.get("LLM_MODEL_NAME")

        logger.info(f"开始调用Deepseek模型生成SQL脚本: model={model}")
        logger.debug(f"输入提示语: {request_data}")

        # 调用Deepseek API生成SQL脚本
        completion = client.chat.completions.create(
            model=model,  # type: ignore[arg-type]
            messages=[
                {
                    "role": "system",
                    "content": "你是一名专业的数据库工程师，专门负责编写高质量的PostgreSQL SQL脚本。"
                    "请严格按照用户提供的需求和表结构信息生成SQL脚本。"
                    "确保生成的SQL语法正确、性能优化，并且能够直接执行。",
                },
                {"role": "user", "content": request_data},
            ],
            temperature=0.1,  # 使用较低的温度确保结果的一致性和准确性
            max_tokens=4096,  # 为SQL脚本提供足够的token空间
            top_p=0.9,
        )

        response_text = completion.choices[0].message.content.strip()  # type: ignore[union-attr]

        logger.info(f"Deepseek模型成功返回SQL脚本，长度: {len(response_text)} 字符")
        logger.debug(f"生成的SQL脚本: {response_text}")

        return response_text

    except Exception as e:
        logger.error(f"Deepseek SQL生成调用失败: {str(e)}")
        raise Exception(f"调用Deepseek模型生成SQL脚本失败: {str(e)}")