langchain_agent.py

import os
import json
import re
import time
import pymysql
from decimal import Decimal
from datetime import datetime, date
from typing import List, Dict, Union, Any, Optional
from pathlib import Path

# 导入 LangChain 核心组件
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

# ================= 1. 配置区域 =================

DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY", "YOUR_API_KEY_HERE") # 请替换为你的 DeepSeek API Key
DEEPSEEK_API_BASE = os.getenv("DEEPSEEK_API_BASE", "https://api.deepseek.com")
DEEPSEEK_MODEL = "deepseek-coder"

DB_CONFIG = {
    'host': '127.0.0.1',
    'user': 'root',
    'password': '',
    'db': 'YOUR_DATABASE_NAME', # 请替换为你的数据库名称
    'port': 9030
}

DATA_DIR = Path("./data")
OUTPUT_FILE = Path("dataset_exe_result.json")

SCHEMA_FILE = DATA_DIR / 'schema.json'
KNOWLEDGE_FILE = DATA_DIR / 'common_knowledge.md'
DATASET_FILE = DATA_DIR / 'final_dataset.json'

# ================= 2. 核心工具类 =================

class DecimalEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, Decimal):
            return int(obj) if obj == obj.to_integral_value() else float(obj)
        elif isinstance(obj, (datetime, date)):
            return obj.isoformat()
        return super().default(obj)

class execute_sql_with_pymysql:
    def __init__(self):
        pass

    def execute_sql_with_pymysql(self, sql_query: str, db_config: Dict[str, Any]) -> List[Dict[str, Any]]:
        conn = None
        result_list = []
        try:
            conn = pymysql.connect(
                host=db_config['host'],
                user=db_config['user'],
                password=db_config.get('password', ''),
                database=db_config['db'],
                port=db_config['port'],
                cursorclass=pymysql.cursors.DictCursor
            )
            with conn.cursor() as cursor:
                cursor.execute(sql_query)
                result_list = cursor.fetchall()
        except Exception as e:
            print(f"Error executing SQL: {e}")
            print(f"Failed SQL Query: {sql_query}")
            raise e 
        finally:
            if conn:
                conn.close()
        return result_list

    def normalize_numbers_in_result(self, result_list: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        def _normalize_value(value: Union[int, float, Decimal, str, None]) -> Union[int, float, str, None]:
            if isinstance(value, (int, float)):
                return int(value) if value == int(value) else round(value, 2)
            if isinstance(value, Decimal):
                return int(value) if value == value.to_integral_value() else round(value, 2)
            return value
        return [{k: _normalize_value(v) for k, v in row.items()} for row in result_list]

# ================= 3. 辅助函数  =================

def load_json_file(filepath: Path):
    try:
        with filepath.open('r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        print(f"Error loading {filepath}: {e}")
        return []

def load_markdown_file(filepath: Path):
    try:
        with filepath.open('r', encoding='utf-8') as f:
            return f.read()
    except Exception as e:
        return ""

def get_pruned_schema_string(table_list: List[str], full_schema: List[Dict]) -> str:
    """
    Schema 剪枝：
    1. 包含表描述 (Table Description)
    2. 包含列描述 (Column Comment) -> 非常重要，枚举值通常在这里！
    """
    if not full_schema: return ""
    schema_str = "--- DATABASE SCHEMA ---\n"
    schema_map = {t.get('table_name'): t for t in full_schema}
    target_tables = table_list if table_list else schema_map.keys()
    
    found = False
    for table_name in target_tables:
        if table_name in schema_map:
            found = True
            t_info = schema_map[table_name]
            # 重点：加入表描述
            t_desc = t_info.get('table_description', '') or t_info.get('description', '')
            
            schema_str += f"Table: {table_name}\n"
            if t_desc:
                schema_str += f"Description: {t_desc}\n"
                
            schema_str += "Columns:\n"
            for col in t_info.get('columns', []):
                c_name = col.get('col', col.get('name', ''))
                c_type = col.get('type', '')
                # 重点：加入列描述/注释
                c_desc = col.get('description', col.get('comment', ''))
                
                schema_str += f"  - {c_name} ({c_type})"
                if c_desc:
                    schema_str += f" : {c_desc}"
                schema_str += "\n"
            schema_str += "\n"
    
    if not found and table_list:
        return get_pruned_schema_string([], full_schema)
    return schema_str

def extract_pure_sql(text: str) -> str:
    match = re.search(r"```sql\s*(.*?)\s*```", text, re.DOTALL | re.IGNORECASE)
    if match: return match.group(1).strip()
    match = re.search(r"```\s*(.*?)\s*```", text, re.DOTALL)
    if match: return match.group(1).strip()
    return text.strip()

# ================= 4. Agent 主逻辑 (针对准确率优化的 Prompt) =================

class StarRocksAgent:
    def __init__(self):
        self.llm = ChatOpenAI(
            model=DEEPSEEK_MODEL,
            openai_api_key=DEEPSEEK_API_KEY,
            openai_api_base=DEEPSEEK_API_BASE,
            temperature=0.0, 
            max_tokens=1024
        )
        self.executor = execute_sql_with_pymysql()
        self.full_schema = load_json_file(SCHEMA_FILE)
        self.common_knowledge = load_markdown_file(KNOWLEDGE_FILE)

    def run_agent(self, question: str, table_list: List[str], item_knowledge: str = '', max_retries=4) -> Dict[str, Any]:
        
        pruned_schema = get_pruned_schema_string(table_list, self.full_schema)
        
        # === 核心修改：Prompt 逻辑重构 ===
        # 1. 移除对日期的强制转换要求，改为“基于类型判断”。
        # 2. 强调从 Schema/Knowledge 中寻找枚举值映射。
        # 3. 提供 One-Shot 示例（可选，这里使用通用指令）。
        
        system_template = f"""你是一个 StarRocks (MySQL 兼容) 数据库专家。
你的任务是编写准确的 SQL 来回答用户问题。

### 1. 数据库 Schema
{pruned_schema}

### 2. 业务知识 (Common Knowledge)
**重要提示：** 请仔细阅读下方的业务知识，其中包含了字段代码（如段位 ID、游戏 ID）的具体映射关系。
{self.common_knowledge}

### 3. 核心规则 (CRITICAL)
1.  **数据值匹配**：
    - 不要猜测枚举值！请从 Schema 注释或业务知识中查找（例如：段位 ID 3-5 代表黑铁，而不是 0-2）。
    - 如果没有找到具体映射，请在注释中说明你假设的值。
    - **数据源完整性：** 如果业务分析（例如：统计所有模式）需要多个数据源（例如 `dws_..._di` 和 `dwd_..._hi`），**考虑使用 `UNION ALL`** 将所有相关数据源合并后再进行聚合计算。
2.  **日期处理 (Date Handling)**：
    - **统一格式**：除非 Schema 描述中明确指出了其他格式（如 'YYYY-MM-DD'），否则所有名为 `dtstatdate`, `statis_date`, `tdbank_imp_date` 等日期字段，**一律转换为 'YYYYMMDD' 格式**。
    - **处理 String 类型**：如果 Schema 中字段类型为 `string`，必须使用单引号，且**去除**日期中的点号、横杠。
      - 错误示例：dtstatdate = '2025-07-17' 或 '2025.7.17'
      - 正确示例：dtstatdate = '20250717'
    - **处理 Bigint/Int 类型**：如果 Schema 中字段类型为 `bigint` 或 `int`，**严禁使用引号**，且必须是纯数字。
      - 错误示例：statis_date = '20250717'
      - 正确示例：statis_date = 20250717
    - **禁止函数**：严禁对表中的分区字段（如 dtstatdate）使用 `date_format`, `to_date` 等函数，这会导致查询极慢或失效。直接比较字符串或数字。
    - **字段选择优先级：**
           **事件时间优先：** 对于涉及用户行为、精确时间范围（例如 **"开服至今"、"活跃期间"**）的业务逻辑过滤，**必须优先使用包含完整时间信息的字段**（例如 `dteventtime` 或 `gamestartmillis`）。
           **分区时间次之：** `tdbank_imp_date`, `dtstatdate` 等字段仅在没有更精确事件时间字段时，或仅用于**分区裁剪**时使用。
3.  **SQL 风格**：
    - 对于‘简单’复杂度的题目，优先使用 `NOT IN`, `EXISTS` 等简单逻辑，避免过度嵌套的 `WITH` 子句，除非非常必要。
    - 对于‘中等’或‘复杂’复杂度的题目，允许适当使用复杂SQL语句。

4.  **输出格式**：
    - **最终输出必须是且只能是**一个完整的、用于回答用户问题的 SQL 代码块。
    - 必须包含 SQL 注释 (`--`) 来解释你的逻辑，特别是你是如何确定过滤值的。
    - 最终输出一段完整的 SQL 代码块。
    - 输出字段与题目要求完全吻合，不要添加、减少或改变输出字段。
    - 输出字段名称、类型及含义严格参考 Schema 中的信息。**使用 Schema 中的英文列名**作为 `SELECT` 的字段，不要使用中文别名（除非题目明确要求输出特定的中文列名）。
5.  **代码调试**：
    - **【调试隔离】** 所有的中间验证（例如 `SELECT COUNT(*)...`）都属于你的内部思考过程，**绝对不能**出现在最终提交的 SQL 代码块中。
    - 严禁在最终 SQL 中包含 `LIMIT 100`（除非题目要求） 或返回非题目要求的中间调试数据。最终输出必须直接回答问题。
6.  **聚合输出**
    - 如果用户问题涉及一个较长的日期范围，且要求的输出字段不包含日期相关字段或其他任何分组字段时，则最终结果必须是针对整个时间范围的单行汇总结果。
    - 示例：如果问题是“统计7月活跃人数。输出：人数”，则只需要 SELECT COUNT(DISTINCT user_id) AS 人数，不需要 GROUP BY dtstatdate。
    - 固定点常量处理 (CRITICAL)：对于像“统计20250226往前180天...”这种固定时间点的滚动指标查询，如果在 SELECT 列表中添加了报告日期（例如 SELECT '20250226' AS dtstatdate, ...），这个日期字段是一个字符串常量，严禁在 SQL 末尾添加 GROUP BY 语句，因为它是常量，不用于分组。

### 4. 思考示例
1.  **日期检查：** 检查日期字段是否符合 'YYYYMMDD' 格式，并核对是否正确使用了引号（`string`用引号，`bigint`不用引号）。
2.  **过滤条件检查：** 暂时注释掉**最严格的过滤条件**（例如：游戏名过滤、渠道过滤、账号类型过滤），以验证数据集合是否为空。
3.  **连接检查：** 确保联表查询中的 `ON` 条件准确无误。
4.  **乱码检查：** **禁止**在 SQL 或注释中提及“Schema缺失”、“付费表不存在”等业务限制。必须生成一个**可执行**的 SQL 语句。
5.  **输出格式检查 ：** 检查是否输出了**唯一的、完整的 SQL 查询**，并确保其格式（`SELECT` 字段）与用户要求一致。在提交最终 SQL 前，**再次确认输出是否只有一个**完整的 SQL 代码块，**不包含任何调试用的中间 `SELECT` 语句**。
......
书写sql时可以添加注释辅助思考。
"""
        
        messages = [
            SystemMessage(content=system_template),
            HumanMessage(content=f"问题: {question}\n\n" + 
                                  (f"***关于此问题的知识库说明***:\n{item_knowledge}\n" 
                                   if item_knowledge else ""))
        ]                                                                                                             

        final_sql = ""
        final_error = None
        last_valid_but_empty_sql = "" 

        display_q = question[:50].replace('\n', ' ')
        print(f"\n=== 处理问题: {display_q}... ===")

        for attempt in range(max_retries + 1):
            print(f"--- 第 {attempt + 1} 次尝试 (Max {max_retries + 1}) ---")
            
            # 1. 生成
            try:
                response = self.llm.invoke(messages)
                ai_content = response.content
                sql_query = extract_pure_sql(ai_content)
                messages.append(AIMessage(content=ai_content))
                
                if "填入" in sql_query or "替换" in sql_query:
                   raise ValueError("Detected placeholder text in SQL.")
                
                print(f"--- 生成 SQL: ---\n{sql_query}\n-------------------")
                
            except Exception as e:
                print(f"LLM 生成错误: {e}")
                messages.append(HumanMessage(content=f"生成出错: {e}。请重新生成。"))
                continue

            # 2. 执行
            try:
                raw_results = self.executor.execute_sql_with_pymysql(sql_query, DB_CONFIG)
                result_count = len(raw_results)
                print(f"--- 原始结果行数: {result_count} ---")
                
                # 3. 结果为空时的【智能反思】
                if result_count == 0:
                    last_valid_but_empty_sql = sql_query
                    
                    if attempt < max_retries:
                        print(">>> 触发【空结果纠错】...")
                        # 针对我们分析出的原因，给出具体的反思提示
                        feedback = """SQL 执行成功但结果为空。请检查以下高频问题：
前提：确保**输出格式检查 ：** 检查是否输出了**唯一的、完整的 SQL 查询**，并确保其格式（`SELECT` 字段）与用户要求一致。
1. **日期类型错误**：Schema 中该日期字段是 INT 还是 String？你是否错误地使用了 `STR_TO_DATE` 或日期格式符（如 '-')？请尝试直接使用纯数字（如 20230101）。
2. **枚举值错误**：你过滤的字段值（如段位 ID、游戏 Code）是否正确？请再次检查 Common Knowledge 或 Schema 注释，不要使用默认的 0,1,2...
3. **复杂逻辑重构 ：** 如果使用了复杂的聚合或窗口函数，若非必要，请尝试简化或将复杂计算拆解。
4.  **业务/过滤条件放宽 ：** 逐步放宽**最不可能出错的过滤条件**（例如：排位/普通模式的 `queueid` 值），以验证数据存在性。
5.  **核心过滤条件：** **严禁**将核心日期范围条件、去重条件、或业务主题（如 FPS）过滤条件完全移除或注释掉。必须找到最合理的代码进行筛选。
6. **JOIN 丢失**：是否 INNER JOIN 条件太严？尝试 LEFT JOIN。
7. **聚合函数错误**：没有理清题目含义，错误使用聚合函数导致读取空值。
请依照上述分析修改 SQL。"""
                        messages.append(HumanMessage(content=feedback))
                        continue
                    else:
                        return {"sql": sql_query, "result": [], "status": "empty"}
                
                # --- NEW: 结果非空时 (result_count > 0) 的【字段空值 (NULL) 检查】 ---
                normalized_results = self.executor.normalize_numbers_in_result(raw_results)
                
                has_null_field = False
                # 检查所有字段值是否包含 None (Python中的null)
                for row in normalized_results:
                    if any(v is None for v in row.values()):
                        has_null_field = True
                        break
                        
                if has_null_field:
                    last_valid_but_empty_sql = sql_query # Store for potential rollback
                    
                    if attempt < max_retries:
                        print(">>> 触发【非空结果中的字段空值纠错】...")
                        feedback = """SQL 执行成功且返回了数据，但**结果集中包含 NULL/空值**。
在最终 SELECT 字段中，**严禁**出现 NULL 值，这通常意味着：
1. **JOIN 失败**：联表查询 (JOIN) 时的 ON 条件不够准确或匹配失败，导致部分字段为空。请检查 JOIN 键和条件。
2. **聚合空值**：聚合函数（如 SUM, AVG）应用于空数据集或过滤条件过于严格，导致计算结果为空。或者题目意义理解错误，错误使用聚合函数。
3. **CASE 语句不完整**：CASE WHEN 表达式缺少 ELSE 分支，或 ELSE 分支返回了 NULL。
请修改 SQL，**确保最终 SELECT 结果的每个字段都有有效值**。"""
                        messages.append(HumanMessage(content=feedback))
                        continue
                    else:
                        # 达到最大重试次数，返回当前结果，不再重试
                        print(">>> 达到最大重试次数，结果包含空值，返回当前结果。")
                        return {"sql": sql_query, "result": normalized_results, "status": "success"}

                # 4. 成功返回 (非空，且无空值)
                return {"sql": sql_query, "result": normalized_results, "status": "success"}

            except Exception as e:
                error_msg = str(e)
                print(f"执行报错: {error_msg}")
                final_error = error_msg
                
                if attempt < max_retries:
                    # 简化的报错反馈
                    hint = "请检查 Schema 字段名是否正确，以及语法是否符合 StarRocks/MySQL 标准。"
                    if "Unknown column" in error_msg:
                        hint = "列名不存在。请仔细核对 Schema，不要编造列名。"
                    
                    feedback = f"执行报错: {error_msg}\n{hint}\n请修正 SQL。"
                    messages.append(HumanMessage(content=feedback))
                else:
                    final_sql = sql_query

        if final_error and last_valid_but_empty_sql:
            print(">>> 回退到空结果版本。")
            return {"sql": last_valid_but_empty_sql, "result": [], "status": "empty"}

        return {"sql": final_sql, "result": [], "status": "error", "error": final_error}

# ================= 5. 统计类 (保持不变) =================

class ExecutionStats:
    def __init__(self):
        self.total = 0
        self.success_with_data = 0
        self.success_empty = 0
        self.error = 0
        self.success_ids = []
        self.empty_ids = []
        self.error_ids = []
    
    def record(self, sql_id: str, status: str, result_count: int):
        self.total += 1
        if status == "success":
            self.success_with_data += 1
            self.success_ids.append(sql_id)
        elif status == "empty":
            self.success_empty += 1
            self.empty_ids.append(sql_id)
        else:
            self.error += 1
            self.error_ids.append(sql_id)
    
    def print_summary(self):
        print("\n" + "=" * 60)
        print("                    执行结果统计报告")
        print("=" * 60)
        print(f"总题目数:                          {self.total}")
        print("-" * 60)
        print(f"✅ 成功且有结果 (result > 0):       {self.success_with_data}  ({self.success_with_data/self.total*100:.1f}%)")
        print(f"⚠️  成功但结果为空 (result = 0):    {self.success_empty}  ({self.success_empty/self.total*100:.1f}%)")
        print(f"❌ 执行报错:                        {self.error}  ({self.error/self.total*100:.1f}%)")
        print("-" * 60)
        if self.error_ids:
            print(f"报错 ID: {', '.join(self.error_ids[:10])}...")
        if self.empty_ids:
            print(f"空结果 ID: {', '.join(self.empty_ids[:10])}...")
        print("")

# ================= 6. 主程序 =================

def main():
    if not OUTPUT_FILE.parent.exists():
        OUTPUT_FILE.parent.mkdir(parents=True, exist_ok=True)

    agent = StarRocksAgent()
    stats = ExecutionStats()
    
    print(f"--- 加载数据集: {DATASET_FILE} ---")
    if not DATASET_FILE.exists():
        print("错误: 数据集不存在")
        return
        
    dataset = load_json_file(DATASET_FILE)
    submission_results = []
    total = len(dataset)
    
    start_time = time.time()

    for i, item in enumerate(dataset):
        sql_id = item.get('sql_id', item.get('id', f'unknown_{i}'))
        question = item.get('question', '')
        table_list = item.get('table_list', [])
        item_knowledge = item.get('knowledge', '')

        print(f"\n>>> Progress {i+1}/{total} | ID: {sql_id}")
        
        result_data = agent.run_agent(question, table_list, item_knowledge)
        
        status = result_data.get('status', 'error')
        result_count = len(result_data.get('result', []))
        stats.record(sql_id, status, result_count)
        
        submission_results.append({
            "sql_id": sql_id,
            "sql": result_data.get('sql', ''),
            "result": result_data.get('result', [])
        })
        
        try:
            with OUTPUT_FILE.open('w', encoding='utf-8') as f:
                json.dump(submission_results, f, indent=4, ensure_ascii=False, cls=DecimalEncoder)
        except Exception as e: pass

    end_time = time.time()
    print(f"\n✅ 处理完成！耗时: {end_time - start_time:.2f}s")
    stats.print_summary()

if __name__ == '__main__':
    main()