MTabishS/Text2SQL

1import re
2import sqlite3
3from groq import Groq
4
5# importing environmental variables and database configuration
6from config import GROQ_API_KEY
7
8# Configure our API key
9client = Groq(
10    api_key=GROQ_API_KEY,
11)
12
13# Function to load Groq Model 
14# and provide SQL query as response
15def get_groq_response(question,prompt):
16    # Preparing the message for Groq Model
17    messages=[
18        {
19            "role":"user",
20            "content":question, # User Query
21        },
22        {
23            "role":"assistant",
24            "content":prompt[0] # initial context or example prompt
25        }
26    ]
27
28    # creating the chat completion
29    chat_completion = client.chat.completions.create(
30        messages=messages,
31        model="llama3-8b-8192"
32    )
33
34    # Extracting and returning the generated SQL Query
35    response_content=chat_completion.choices[0].message.content
36    print('Response Object:',response_content)
37
38    #Using regex to extract SQL query between ```sql and ```
39    sql_query=""
40    match=re.search(r'```sql\n+(.*?)\n+```',response_content,re.DOTALL)
41
42    if match:
43        sql_query=match.group(1).strip()
44    else:
45        #print('Could not find SQL in the expected format.')
46        #sql_query=response_content.strip()
47        raise ValueError('No Valid SQL query found in response')
48    
49    print('Generated SQL Query:',sql_query)
50    return sql_query
51
52
53# Function to retrieve query from SQL database
54def read_sql_query(sql,db):
55    try:
56        conn=sqlite3.connect(db)
57        cur=conn.cursor()
58        cur.execute(sql)
59        rows=cur.fetchall()
60
61        # if query return no rows
62        if not rows:
63            print('No data found for the query')
64            return[]
65    
66        # converting rows to list of dictionaries
67        columns=[description[0] for description in cur.description]
68        data = [dict(zip(columns,row)) for row in rows]
69
70        conn.commit()
71        conn.close()
72        
73        return data
74        
75    except sqlite3.OperationalError as e:
76        print(f'OperationalError: {e}')
77        return[]
78    except Exception as e:
79        print(f'An Error Occurred: {e}')
80        return[]
81    
82
83# Defining Prompt to make Model understand our requirements
84prompt = [
85    """
86    You are an expert in converting English questions to SQL queries!
87
88    The SQL database has six tables as follows:
89
90    Apartment_Bookings: (apt_booking_id, apt_id, guest_id, booking_status_code, booking_start_date, booking_end_date)
91    Apartment_Buildings: (building_id, building_short_name, building_full_name, building_description, building_address, building_manager, building_phone)
92    Apartment_Facilities: (apt_id, facility_code)
93    Apartments: (apt_id, building_id, apt_type_code, apt_number, bathroom_count, bedroom_count, room_count)
94    Guests: (guest_id, gender_code, guest_first_name, guest_last_name, date_of_birth)
95    View_Unit_Status: (apt_id, apt_booking_id, statust_date, available_yn)
96    When converting English questions to SQL queries, consider the following examples:
97
98    Example 1: Simple Select Query
99
100    Question: "Show me all the details of apartments."
101    SQL: SELECT * FROM Apartments;
102    Example 2: Conditional Select Query
103
104    Question: "List the names of all buildings managed by 'John Doe'."
105    SQL: SELECT building_full_name FROM Apartment_Buildings WHERE building_manager = 'John Doe';
106    Example 3: Join Query
107
108    Question: "Find the booking details for the apartment with apartment ID 'A101'."
109    SQL: SELECT * FROM Apartment_Bookings WHERE apt_id = 'A101';
110    Example 4: Aggregation Query
111
112    Question: "How many apartments are in each building?"
113    SQL: SELECT building_id, COUNT(*) as apartment_count FROM Apartments GROUP BY building_id;
114    Example 5: Complex Join Query
115
116    Question: "List the guest names and their booking statuses for all bookings starting after '2024-07-01'."
117    SQL: SELECT g.guest_first_name, g.guest_last_name, ab.booking_status_code FROM Guests g JOIN Apartment_Bookings ab ON g.guest_id = ab.guest_id WHERE ab.booking_start_date > '2024-07-01';
118    Example 6: Filtering with Multiple Conditions
119
120    Question: "Show the available units as of '2024-07-15'."
121    SQL: SELECT * FROM View_Unit_Status WHERE statust_date = '2024-07-15' AND available_yn = 'Y';
122    When converting English questions to SQL:
123
124    Ensure that table and column names are used correctly.
125    Join tables where necessary to retrieve information across multiple tables.
126    Apply appropriate filtering conditions (WHERE clauses) based on the query requirements.
127    Aggregate data when needed (e.g., using COUNT, SUM, AVG, etc.).
128    Order the results if specified (e.g., using ORDER BY).
129
130    For all responses, provide only a valid SQL query without any extra explanation or comments. Wrap the SQL query in triple backticks ```sql and end with ```.
131
132    """
133]