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

import os
import requests
import re
import json
from flask import current_app

class DDLParser:
    def __init__(self, api_key=None):
        """
        初始化DDL解析器
        
        参数:
            api_key: LLM API密钥，如果未提供，将从应用配置或环境变量中获取
        """
        # 如果在Flask应用上下文中，则从应用配置获取参数
       
        self.api_key = api_key or current_app.config.get('LLM_API_KEY')
        self.base_url = current_app.config.get('LLM_BASE_URL')
        self.model_name = current_app.config.get('LLM_MODEL_NAME')
        
        
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

    def parse_ddl(self, sql_content):
        """
        解析DDL语句，返回标准化的结构
        
        参数:
            sql_content: 要解析的DDL语句
            
        返回:
            解析结果的JSON对象
        """
        prompt = self._optimize_prompt()
        payload = {
            "model": self.model_name,
            "messages": [
                {
                    "role": "system",
                    "content": "你是一个专业的数据库分析专家，擅长解析SQL DDL语句并提取表结构信息。"
                },
                {
                    "role": "user", 
                    "content": f"{prompt}\n\n{sql_content}"
                }
            ]
        }
        
        try:
            response = requests.post(
                f"{self.base_url}/chat/completions",
                headers=self.headers,
                json=payload,
                timeout=30
            )
            response.raise_for_status()
            
            result = response.json()
            
            if "choices" in result and len(result["choices"]) > 0:
                content = result["choices"][0]["message"]["content"]
                
                try:
                    json_match = re.search(r'```json\s*([\s\S]*?)\s*```', content)
                    if json_match:
                        json_content = json_match.group(1)
                    else:
                        json_content = content
                        
                    parsed_result = json.loads(json_content)
                    return parsed_result
                except json.JSONDecodeError as e:
                    return {
                        "code": 500,
                        "message": f"无法解析返回的JSON: {str(e)}",
                        "original_response": content
                    }
            
            return {
                "code": 500,
                "message": "无法获取有效响应",
                "original_response": result
            }
            
        except requests.RequestException as e:
            return {
                "code": 500,
                "message": f"API请求失败: {str(e)}"
            }

    def _optimize_prompt(self):
        """返回优化后的提示词模板"""
        return """
请解析以下DDL建表语句，并按照指定的JSON格式返回结果：

规则说明:
1. 从DDL语句中识别所有表名，并在data对象中为每个表创建条目,表名请使用小写，可能会有多个表。
2. 对于每个表，提取所有字段信息，包括名称、数据类型和注释。
   - 中文表名中不要出现标点符号
3. 字段中文名称(name)的确定规则：
   - 如有COMMENT注释，直接使用注释内容
   - 如无注释但字段名有明确含义，将英文名翻译为中文
   - 如字段名是无意义的拼音缩写，则name为空字符串
   - 字段名中不要出现逗号，以及"主键"、"外键"、"索引"等字样
4. 数据库连接串处理：
   - 将连接串识别后并拆解为：主机名/IP地址、端口、数据库名称、用户名、密码。
   - 根据连接串格式识别数据库类型，数据库类型请使用小写，参考例子，如 mysql/postgresql/sqlserver/oracle/db2/sybase
   - data_source.en_name格式为: "{数据库名称}_{hostname或ip地址}_{端口}_{数据库用户名}",如某个元素无法识别，则跳过不添加.
   - data_source.name留空.
   - 无法确定数据库类型时，type设为"unknown"
   - 如果从ddl中没有识别到数据库连接串，则json不返回"data_source"标签
   - 除了database,password,username,en_name,host,port,type,name 之外，连接串的其它字段放在param属性中。
5. 参考格式如下：
{
    "users": { //表名
        "name": "用户表", //表的中文名，来自于COMMENT注释或LLM翻译，如果无法确定，则name为空字符串
        "schema": "public",
        "meta": [{
                "en_name": "id",
                "data_type": "integer",
                "name": "用户ID"
            },
            {
                "en_name": "username",
                "data_type": "varchar",
                "name": "用户名"
            }
        ]
    },
    "data_source": [{
        "en_name": "mydatabase_10.52.31.104_5432_myuser", //{数据库名称}_{hostname或ip地址}_{端口}_{数据库用户名}
        "name": "", //如果没有注释，这里留空
        "type": "postgresql",
        "host": "10.52.31.104",
        "port": 5432,
        "database": "mydatabase",
        "username": "myuser",
        "password": "mypassword",
        "param": "useUnicode=true&characterEncoding=utf8&serverTimezone=UTC"
    }]
}

请仅返回JSON格式结果，不要包含任何其他解释文字。
"""