agents/sql_agent.py

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)