File size: 1,413 Bytes
1914b78
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
## as i fine tune model on following message set 





def generate_message_template(User_query ,table_schema):

    messages = [
        {"role":"system" , 
        "content":"You are a SQL generator. Only output SQL query. no explaination, no markdown."} , 
        
         {
        "role":"user" , 
        "content":f"""{User_query},
                    Schema: {table_schema}"""     
         }
    ] 

    return messages

def generate_refine_template(user_query: str, table_schema, previous_sql: str, stage: int):
    """Stage 2 & 3: Refine previous SQL"""
    stage_instructions = {
        2: "Review the SQL query below for logical errors, wrong joins, or missing conditions. Output only the corrected SQL.",
        3: "Final check: ensure the SQL is optimized, uses correct aggregations, and fully answers the user question. Output only the final SQL."
    }

    messages = [
        {
            "role": "system",
            "content": "You are a SQL expert. Only output SQL query. No explanation, no markdown."
        },
        {
            "role": "user",
            "content": f"{user_query},\n                    Schema: {table_schema}"
        },
        {
            "role": "assistant",
            "content": previous_sql   # previous stage output
        },
        {
            "role": "user",
            "content": stage_instructions[stage]
        }
    ]
    return messages