sql_engine.py

import os
import re
import pandas as pd
from urllib.parse import quote_plus
from dotenv import load_dotenv
from langchain_community.utilities import SQLDatabase
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_google_genai import ChatGoogleGenerativeAI

load_dotenv()

def get_db_connection():
    """Initializes and returns the SQL database connection."""
    host = os.getenv("DB_HOST")
    port = os.getenv("DB_PORT")
    username = os.getenv("DB_USER")
    password = quote_plus(os.getenv("DB_PASSWORD", ""))
    database_schema = os.getenv("DB_NAME")
    
    mysql_uri = f"mysql+pymysql://{username}:{password}@{host}:{port}/{database_schema}"
    return SQLDatabase.from_uri(mysql_uri, sample_rows_in_table_info=2)

def get_sql_chain(db):
    """Creates the LangChain SQL generation chain."""
    template = """Based on the table schema below, write a SQL query that would answer the user's question:
Remember : Only provide me the sql query dont include anything else.
           Provide me sql query in a single line dont add line breaks.
Table Schema:
{schema}

Question: {question}
SQL Query:
"""
    prompt = ChatPromptTemplate.from_template(template)
    
    llm = ChatGoogleGenerativeAI(
        model="gemini-2.5-flash",
        api_key=os.getenv("GOOGLE_API_KEY")
    )
    
    def get_schema(_):
        return db.get_table_info()
        
    sql_chain = (
        RunnablePassthrough.assign(schema=get_schema)
        | prompt
        | llm.bind(stop=["\nSQLResult:"])
        | StrOutputParser()
    )
    return sql_chain

def run_text_to_sql(question, db, sql_chain):
    """Processes the natural language query, extracts SQL, and executes it via Pandas."""
    raw_response = sql_chain.invoke({"question": question})
    
    # Clean output using your regex pattern
    match = re.search(r"```sql\s*(.*?)\s*```", raw_response, re.DOTALL | re.IGNORECASE)
    query = match.group(1).strip() if match else raw_response.strip()
    
    try:
        # ADVANCEMENT: Read the SQL query directly into a Pandas DataFrame
        df_result = pd.read_sql_query(query, db._engine)
        return query, df_result, None
    except Exception as e:
        return query, None, str(e)
def generate_summary(question, result_df):
    """
    Generates an AI summary of the SQL output.
    """

    if result_df is None or result_df.empty:
        return "No data was returned for this query."

    llm = ChatGoogleGenerativeAI(
        model="gemini-2.5-flash",
        api_key=os.getenv("GOOGLE_API_KEY")
    )

    prompt = f"""
You are an expert business analyst.

User Question:
{question}

Query Result:
{result_df.head(20).to_string(index=False)}

Write a concise summary (5-7 sentences).

Rules:
- Do not mention SQL.
- Explain the important insight.
- Mention important trends if visible.
- Keep the language simple and professional.
"""

    return llm.invoke(prompt).content