Prompt Engineering在SQL生成任务中的调优实录：从42%到87%准确率

# 实验1：零样本模板
zero_shot_prompt = """### Instruction:
Translate the following natural language query into SQL.

### Input:
{question}

### Schema:
{schema}

### SQL:
"""

# 实验2：指令增强模板（加了任务描述和输出格式）
instruct_prompt = """### Instruction:
You are an expert SQL developer. Given the database schema and a question, generate a SQLite SQL query that answers the question. 
- Only output the SQL statement, no explanations.
- Use only columns and tables from the given schema.
- If the question is ambiguous, make a reasonable assumption and note it in a comment.

### Schema:
{schema}

### Question:
{question}

### SQL:
"""

few_shot_cot_prompt = """### Instruction:
You are an expert SQL developer. For each question, first reason about the steps to solve it, then output the SQL.

### Examples:
Q: How many students have a GPA above 3.5?
A: 
1. Identify the table: student.
2. Filter condition: GPA > 3.5.
3. Count rows.
SQL: SELECT COUNT(*) FROM student WHERE gpa > 3.5;

Q: List the names of departments that have more than 2 courses.
A:
1. Join department and course on dept_id.
2. Group by department name.
3. Having count(course_id) > 2.
SQL: SELECT d.name FROM department d JOIN course c ON d.dept_id = c.dept_id GROUP BY d.name HAVING COUNT(c.course_id) > 2;

### Schema:
{schema}

### Question:
{question}

### Reasoning:
"""

# 注意：这里为了让模型输出推理过程，需要修改采样参数
sampling_params = SamplingParams(temperature=0.1, max_tokens=1024, stop=["SQL:"])

error_fix_prompt = """### Instruction:
... (同上)

### Examples with errors:
Q: Which department has the highest average salary?
Wrong SQL: SELECT dept_name FROM department ORDER BY avg(salary) LIMIT 1;
(Error: salary is in employee table, need join)
Correct SQL: SELECT d.dept_name FROM department d JOIN employee e ON d.dept_id = e.dept_id GROUP BY d.dept_name ORDER BY AVG(e.salary) DESC LIMIT 1;

### Schema:
{schema}

### Question:
{question}

### Reasoning:
"""

structured_prompt = """### Instruction:
You are an expert SQLite developer. Given the schema and question, output a JSON object with:
- "reasoning": a brief step-by-step plan (max 3 steps)
- "sql": the final SQL query

Rules:
- Only use tables/columns in the schema.
- SQL must be valid SQLite syntax.
- Use proper JOIN conditions.

### Schema (relevant tables only):
{schema_relevant}

### Question:
{question}

### Output (JSON):
"""

# 解析JSON并提取SQL
import json
def extract_sql(llm_output):
    try:
        data = json.loads(llm_output)
        return data["sql"]
    except:
        # fallback: 正则提取
        return re.search(r'"sql"\s*:\s*"([^"]+)"', llm_output).group(1)

sampling_params = SamplingParams(temperature=0.3, max_tokens=1024, n=5)

# 投票逻辑：执行所有候选SQL，选择结果集出现次数最多的
from collections import Counter
def self_consistency_sql(candidates, db_path):
    results = []
    for sql in candidates:
        try:
            conn = sqlite3.connect(db_path)
            result = conn.execute(sql).fetchall()
            results.append((sql, str(result)))
            conn.close()
        except:
            results.append((sql, "__ERROR__"))
    counter = Counter([r[1] for r in results if r[1] != "__ERROR__"])
    if not counter:
        return None
    most_common = counter.most_common(1)[0][0]
    return [r[0] for r in results if r[1] == most_common][0]

### System:
You are an expert SQLite query generator. Follow these rules strictly:
1. Only use tables and columns listed in the schema.
2. If the question requires aggregation (avg, max, count), always use GROUP BY on the non-aggregated columns.
3. Use LEFT JOIN only when the question explicitly asks for "all" records from one side.
4. Always use DISTINCT when joining to avoid duplicate rows, unless the question implies duplicates should be kept.

### Schema (relevant tables):
{schema_relevant}

### Previous examples of correct SQL:
{dynamic_few_shot}

### Question:
{question}

### Output JSON format:
{"reasoning": "step-by-step plan", "sql": "the SQL query"}