File size: 2,188 Bytes
e0560d9 | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | import os
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_core.messages import HumanMessage, SystemMessage
from sqlalchemy import create_engine, text
import pandas as pd
from prompts.sql_prompt import SYSTEM_PROMPT
from utils.sql_utils import clean_sql_query
load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
llm = ChatGroq(
groq_api_key=GROQ_API_KEY,
model_name="llama-3.3-70b-versatile"
)
engine = create_engine("sqlite:///database/ecommerce.db")
def generate_sql(user_question):
messages = [
SystemMessage(content=SYSTEM_PROMPT),
HumanMessage(content=user_question)
]
response = llm.invoke(messages)
sql_query = response.content.strip()
cleaned_query = clean_sql_query(sql_query)
return cleaned_query
# =====================================
# Execute SQL
# =====================================
def execute_sql(sql_query):
try:
with engine.connect() as connection:
result = connection.execute(text(sql_query))
rows = result.fetchall()
columns = result.keys()
df = pd.DataFrame(rows, columns=columns)
return df
except Exception as e:
return f"SQL ERROR: {str(e)}"
if __name__ == "__main__":
question = input("Ask your question: ")
# Generate SQL
sql_query = generate_sql(question)
print("\nGenerated SQL:\n")
print(sql_query)
print("\nExecuting Query...\n")
# Execute SQL
result = execute_sql(sql_query)
# Retry if SQL fails
if isinstance(result, str) and "SQL ERROR" in result:
print("\nRetrying with error feedback...\n")
retry_prompt = f"""
The following SQL query failed:
{sql_query}
Error:
{result}
Generate corrected SQLite SQL only.
"""
corrected_sql = generate_sql(retry_prompt)
print("\nCorrected SQL:\n")
print(corrected_sql)
result = execute_sql(corrected_sql)
print("\nFinal Result:\n")
print(result) |