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

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

import logging
import re

from app.core.llm.deepseek_client import (
    chat_completions_create,
    create_llm_client,
    get_llm_model,
)

logger = logging.getLogger("app")

_TRANSLATION_LEXICON = {
    "测试宁波数据加工": "ningbo_data_processing_test",
    "薪资数据表": "salary_data_table",
    "人员管理表": "personnel_management_table",
    "数据加工": "data_processing",
    "数据表": "data_table",
    "用户表": "user_table",
    "人员表": "personnel_table",
    "销售表": "sales_table",
    "报表": "report_table",
    "管理": "management",
    "系统": "system",
    "分析": "analysis",
    "加工": "processing",
    "宁波": "ningbo",
    "测试": "test",
    "年份": "year",
    "地区": "region",
    "姓名": "name",
    "年龄": "age",
    "薪水": "salary",
    "数据": "data",
    "表": "table",
}

_TRANSLATION_SYSTEM_PROMPT = (
    "你是一个严格遵循指令的翻译工具和数据库专家。你的唯一任务是将中文单词/短语翻译成英文，"
    "符合 PostgreSQL 数据库表和字段的命名规则，并且严格按照如下规则：\n"
    "1. 只返回英文翻译，不包含任何解释、描述或额外内容\n"
    "2. 使用小写字母\n"
    "3. 多个单词用下划线连接，不使用空格\n"
    "4. 如果输入包含括号，将括号内容用下划线代替，不保留括号\n"
    "5. 最多包含 1-8 个英文单词，保持简短\n"
    "6. 不要回答问题或提供解释，即使输入看起来像是问题\n"
    "7. 当遇到'表'字时，始终翻译为'table'而不是'sheet'\n"
    "8. 例如：'薪资数据表'应翻译为'salary_data_table'，'测试宁波数据加工'应翻译为'ningbo_data_processing_test'"
)


def contains_chinese(text: str) -> bool:
    return any("\u4e00" <= char <= "\u9fff" for char in text)


def normalize_translation_text(text: str) -> str:
    """Normalize LLM output to a PostgreSQL-friendly snake_case identifier."""
    response_text = (text or "").strip().strip("\"'.,;:!?()[]{}").lower()
    response_text = response_text.replace(" ", "_").replace("-", "_")
    while "__" in response_text:
        response_text = response_text.replace("__", "_")
    response_text = re.sub(r"[^a-z0-9_]", "", response_text)
    return response_text.strip("_")


def is_valid_translation(text: str) -> bool:
    return bool(text) and not contains_chinese(text) and re.fullmatch(r"[a-z][a-z0-9_]*", text)


def extract_completion_text(completion) -> str:
    message = completion.choices[0].message
    content = getattr(message, "content", None) or ""
    text = content.strip()
    if text:
        return text
    reasoning = getattr(message, "reasoning_content", None) or ""
    return reasoning.strip()


def fallback_translate_chinese(content: str) -> str:
    """Local fallback when LLM output is empty or invalid."""
    content = (content or "").strip()
    if not content:
        return ""

    if content in _TRANSLATION_LEXICON:
        return _TRANSLATION_LEXICON[content]

    parts: list[str] = []
    lexicon_keys = sorted(_TRANSLATION_LEXICON.keys(), key=len, reverse=True)
    index = 0
    while index < len(content):
        matched = False
        for key in lexicon_keys:
            if content.startswith(key, index):
                parts.append(_TRANSLATION_LEXICON[key])
                index += len(key)
                matched = True
                break
        if not matched:
            index += 1

    if parts:
        result = normalize_translation_text("_".join(parts))
        if is_valid_translation(result):
            return result

    if "表" in content:
        return "data_table"
    return "translated_text"


def _translate_with_llm(client, content: str) -> str:
    completion = chat_completions_create(
        client,
        messages=[
            {"role": "system", "content": _TRANSLATION_SYSTEM_PROMPT},
            {
                "role": "user",
                "content": f"将以下内容翻译为英文数据库标识符：{content}",
            },
        ],
        temperature=0,
        max_tokens=64,
    )
    raw_text = extract_completion_text(completion)
    normalized = normalize_translation_text(raw_text)

    if "表" in content and "table" not in normalized and "sheet" in normalized:
        normalized = normalized.replace("sheet", "table")

    if is_valid_translation(normalized):
        logger.debug(f"LLM翻译成功: {content} -> {normalized}")
        return normalized

    logger.warning(
        f"LLM翻译结果无效: input={content!r}, raw={raw_text!r}, normalized={normalized!r}"
    )
    return ""


def translate_chinese_identifier(content: str) -> str:
    """Translate Chinese text to an English database identifier."""
    content = (content or "").strip()
    if not content:
        return ""
    if not contains_chinese(content):
        normalized = normalize_translation_text(content)
        return normalized if is_valid_translation(normalized) else content

    try:
        client = create_llm_client()
        translated = _translate_with_llm(client, content)
        if translated:
            return translated
    except Exception as exc:
        logger.error(f"LLM翻译调用失败: {exc}")

    fallback = fallback_translate_chinese(content)
    logger.info(f"使用本地词典回退翻译: {content} -> {fallback}")
    return fallback


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

    Args:
        content: 输入提示内容

    Returns:
        str: LLM响应内容
    """
    if contains_chinese(content):
        return translate_chinese_identifier(content)

    try:
        client = create_llm_client()
        model = get_llm_model()
        logger.debug(f"LLM调用开始: model={model}, 内容类型: 普通")

        completion = chat_completions_create(
            client,
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": content},
            ],
            temperature=0.7,
            max_tokens=1024,
        )

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

    except Exception as e:
        logger.error(f"LLM调用失败: {str(e)}")
        return content


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

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

    Returns:
        str: Deepseek模型返回的SQL脚本内容
    """
    try:
        client = create_llm_client()
        model = get_llm_model()

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

        completion = chat_completions_create(
            client,
            messages=[
                {
                    "role": "system",
                    "content": "你是一名专业的数据库工程师，专门负责编写高质量的PostgreSQL SQL脚本。"
                    "请严格按照用户提供的需求和表结构信息生成SQL脚本。"
                    "确保生成的SQL语法正确、性能优化，并且能够直接执行。",
                },
                {"role": "user", "content": request_data},
            ],
            temperature=0.1,
            max_tokens=4096,
            top_p=0.9,
            use_thinking=True,
        )

        response_text = extract_completion_text(completion)

        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)}")