企业级NL2SQL的高危拦截与自愈：基于腾讯混元+VectorDB的智能数据分析Agent实践

import sqlglot
from sqlglot import exp

class SQLSafetyChecker:
    FORBIDDEN_TOKENS = ['DROP', 'DELETE', 'UPDATE', 'INSERT', 'ALTER', 'CREATE', 'TRUNCATE', 'GRANT', 'REVOKE']
    
    @staticmethod
    def validate(sql: str) -> tuple[bool, str]:
        """
        返回 (是否安全, 拒绝原因)
        """
        try:
            parsed = sqlglot.parse_one(sql, dialect="mysql")
        except Exception as e:
            return False, f"SQL语法解析失败: {str(e)}"

        # 1. 检查顶层操作类型
        if not isinstance(parsed, exp.Select):
            return False, f"仅支持SELECT查询，禁止执行{parsed.key}操作"

        # 2. 深度遍历，禁止多表UPDATE或子查询中的修改操作（部分数据库支持with update）
        for node in parsed.walk():
            if isinstance(node, exp.Update) or isinstance(node, exp.Delete):
                return False, f"检测到高危操作: {node.key}"
            
            # 检查表名是否包含敏感库（如mysql系统库）
            if isinstance(node, exp.Table):
                if node.name.lower() in ['mysql', 'information_schema', 'performance_schema']:
                    return False, "禁止查询系统元数据库"

        # 3. 限制返回行数，强制添加LIMIT防止全表扫描拖垮数据库
        if not parsed.args.get('limit'):
            # 如果SQL没有LIMIT，我们自动追加（但不修改原SQL，在后续执行层强制注入）
            pass 
            
        return True, "安全校验通过"

# 测试用例
print(SQLSafetyChecker.validate("SELECT * FROM users"))  # (True, '安全校验通过')
print(SQLSafetyChecker.validate("DELETE FROM users WHERE id=1"))  # (False, '仅支持SELECT...')

import tcvectordb
from tcvectordb.model.enum import FieldType, IndexType, MetricType
from tencentcloud.common import credential
from sentence_transformers import SentenceTransformer  # 可使用本地BGE模型

# 初始化客户端（内网地址）
client = tcvectordb.VectorDBClient(url="http://your-vip.vectordb.tencentcloudapi.com", 
                                   username="root", key="your-key", timeout=30)

# 创建Collection（如果不存在）
collection = client.create_collection(
    database_name="nl2sql_meta",
    collection_name="field_embeddings",
    shard=1,
    replicas=1,
    description="数据字典向量索引",
    index_spec=[
        {"field_name": "id", "field_type": FieldType.String, "index_type": IndexType.PRIMARY_KEY},
        {"field_name": "vector", "field_type": FieldType.Vector, "index_type": IndexType.HNSW,
         "dimension": 768, "metric_type": MetricType.COSINE},
        {"field_name": "field_name", "field_type": FieldType.String},
        {"field_name": "table_name", "field_type": FieldType.String},
        {"field_name": "comment", "field_type": FieldType.String}
    ]
)

# 入库示例（实际需遍历所有表的INFORMATION_SCHEMA）
def index_field(field_name, table_name, comment, embedding_vector):
    doc = {
        "id": f"{table_name}_{field_name}",
        "vector": embedding_vector,
        "field_name": field_name,
        "table_name": table_name,
        "comment": comment
    }
    collection.upsert(documents=[doc])

def retrieve_relevant_schema(question: str, top_k=5):
    # 调用混元Embedding接口或本地BGE
    # 这里假设已有 embed(text) 函数
    q_vector = embed(question)
    
    res = collection.search(
        vectors=[q_vector],
        limit=top_k,
        output_fields=["field_name", "table_name", "comment"]
    )
    
    # 组装成简洁的DDL描述
    schema_lines = []
    for doc in res[0]:
        schema_lines.append(f"表 {doc['table_name']} 字段 {doc['field_name']} 含义: {doc['comment']}")
    return "\n".join(schema_lines)

from tencentcloud.common import credential
from tencentcloud.hunyuan.v20230901 import hunyuan_client, models
import json

def call_hunyuan(prompt: str) -> str:
    cred = credential.Credential("SECRET_ID", "SECRET_KEY")
    client = hunyuan_client.HunyuanClient(cred, "ap-guangzhou")
    
    req = models.ChatCompletionsRequest()
    req.Model = "hunyuan-pro"
    req.Messages = [{"Role": "user", "Content": prompt}]
    req.Temperature = 0.1  # 降低随机性保证SQL确定性
    
    resp = client.ChatCompletions(req)
    return json.loads(resp.to_json_string())["Choices"][0]["Message"]["Content"]

def execute_with_self_healing(question, df_schema, max_retries=2):
    # 构建初始Prompt
    base_prompt = f"""
你是一个严谨的Hive/MySQL专家。根据表结构生成标准SQL。
要求：
1. 只生成SELECT语句。
2. 不要使用中文别名，不要包含注释。
3. 若涉及时间分区，必须带分区过滤条件（若字段存在）。

表结构信息：
{df_schema}

用户问题：{question}
SQL:
"""
    sql = call_hunyuan(base_prompt)
    
    for attempt in range(max_retries + 1):
        # 1. 安全校验
        is_safe, reason = SQLSafetyChecker.validate(sql)
        if not is_safe:
            # 让混元反思：这是高危SQL，重写
            fix_prompt = f"之前生成的SQL被安全系统拦截，原因：{reason}。请重新生成完全不同的、只包含SELECT的查询。原问题：{question}\n新SQL:"
            sql = call_hunyuan(fix_prompt)
            continue
        
        # 2. 注入强制LIMIT（若没有）
        import re
        if not re.search(r'\bLIMIT\s+\d+', sql, re.IGNORECASE):
            sql = sql.rstrip(';') + " LIMIT 100;"
        
        # 3. 执行（调用腾讯云CDW API或SCF，这里模拟执行）
        exec_result = execute_in_sandbox(sql)  # 返回 {status, data, error_msg}
        
        if exec_result['status'] == 'success':
            return {"sql": sql, "data": exec_result['data']}
        else:
            # 自愈：把错误信息喂给混元
            error_prompt = f"""
之前生成的SQL执行报错：
SQL: {sql}
错误信息: {exec_result['error_msg']}
请根据错误修正SQL，重新输出修正后的完整SQL。
用户问题：{question}
修正后的SQL:
"""
            sql = call_hunyuan(error_prompt)
    
    return {"sql": sql, "error": "超过最大重试次数，执行失败"}