Text2SQL智能体：把“人话”翻译成数据库语言，AI真的做到了

class Text2SQLAgent:
    def __init__(self, llm, schema_repository, sql_executor):
        self.llm = llm                 # 大语言模型（推理核心）
        self.schema = schema_repository # 数据库元数据（表结构、注释）
        self.executor = sql_executor    # SQL执行器（可返回错误信息）
        self.history = []               # 多轮交互记录
    
    def ask(self, natural_language_question: str) -> dict:
        # 1. 意图理解与Schema匹配
        context = self._understand_question(natural_language_question)
        
        # 2. SQL生成
        sql = self._generate_sql(natural_language_question, context)
        
        # 3. 执行与自修正
        result, error = self._execute_with_retry(sql, max_retries=3)
        
        # 4. 结果解释
        return self._format_result(result, natural_language_question)

def _understand_question(self, question: str) -> dict:
    # 先从向量数据库中检索最相关的表和字段
    # 而不是把整个数据库的Schema全塞进提示词（会超过上下文窗口）
    relevant_tables = self.schema.retrieve_relevant(question, top_k=5)
    
    # 让模型把用户问句中的实体映射到具体的表和字段
    linking_prompt = f"""
    数据库中有以下表和字段：
    {self.schema.describe(relevant_tables)}
    
    用户的问题：{question}
    
    请判断：
    1. 用户主要关注哪些表？
    2. 用户提到的“客户”、“订单”、“金额”分别对应哪些字段？
    3. 是否需要关联多张表？如果需要，关联条件是什么？
    
    以JSON格式输出。
    """
    linking_result = self.llm.generate(linking_prompt)
    return json.loads(linking_result)

表名：orders（订单表）
字段：
- order_id (bigint) ：订单编号
- user_id (bigint) ：客户ID，关联users表的user_id
- status (varchar) ：订单状态（pending/paid/shipped/cancelled）
- total_amount (decimal) ：订单总金额
- created_at (timestamp) ：下单时间

def _generate_sql(self, question: str, context: dict) -> str:
    prompt = f"""
    你是一个SQL专家，请根据用户的自然语言问题生成MySQL查询语句。
    
    【数据库Schema】
    {context['schema_description']}
    
    【已识别的字段映射】
    {context['linking_result']}
    
    【业务规则】
    - 所有金额单位均为分（存储为整数）
    - 查询订单时，默认排除状态为 'cancelled' 的订单
    - 涉及时间范围时，如果没有明确指定，默认查询最近90天
    - 排序规则：按创建时间降序排列
    
    【用户问题】
    {question}
    
    【输出要求】
    只输出SQL语句本身，不要添加任何解释文字，不要加markdown代码块。
    """
    
    sql_candidate = self.llm.generate(prompt)
    # 清理潜在的markdown包裹
    return clean_sql(sql_candidate)

def _execute_with_retry(self, sql: str, max_retries: int = 3):
    for attempt in range(max_retries):
        try:
            result = self.executor.execute(sql)
            return result, None
        except Exception as e:
            error_msg = str(e)
            print(f"SQL执行失败（第{attempt+1}次尝试）：{error_msg}")
            
            if attempt == max_retries - 1:
                return None, error_msg
            
            # 让模型根据错误信息修正SQL
            fix_prompt = f"""
            原SQL语句：
            {sql}
            
            执行时出现错误：
            {error_msg}
            
            请修正上述SQL语句，只输出修正后的SQL。
            """
            sql = self.llm.generate(fix_prompt)
    
    return None, "经过多次尝试仍无法生成可执行的SQL"

class Text2SQLSession:
    def __init__(self, session_id: str, agent: Text2SQLAgent):
        self.id = session_id
        self.agent = agent
        self.history = []        # 对话历史
        self.last_sql = None     # 上一轮生成的SQL
        self.last_context = None # 上一轮的Schema上下文
    
    def ask(self, question: str):
        # 如果是连续对话，把历史对话和上一次的SQL作为额外上下文
        enriched_question = self._enrich_with_history(question)
        return self.agent.ask(enriched_question)

def _safety_check(sql: str) -> bool:
    dangerous_keywords = ["DROP", "TRUNCATE", "ALTER", "DELETE"]
    if any(kw in sql.upper() for kw in dangerous_keywords):
        raise SecurityError("检测到危险操作，已拦截")
    if "WHERE" not in sql.upper() and "SELECT" in sql.upper():
        # 无条件的全表查询，需要二次确认
        return False
    return True