app.py

import streamlit as st
import os
import pandas as pd
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from sqlalchemy import create_engine, text, inspect

# Load environment variables
load_dotenv()

# SQL Server Configuration
SERVER_NAME = st.secrets.get("HOST", "SHRABANI")
DATABASE_NAME = st.secrets.get("DATABASE", "demo")
USERNAME = st.secrets.get("USER", "python_user")
PASSWORD = st.secrets.get("PWD", "test")
CSV_PATH = st.secrets.get("path", "")

# Initialize session state for database connection and schema
if 'engine' not in st.session_state:
    st.session_state.engine = None
if 'schema_info' not in st.session_state:
    st.session_state.schema_info = None

@st.cache_resource
def get_sql_engine():
    """Create SQL Server connection"""
    try:
        # Connection string for SQL Server
        conn_string = f"mssql+pyodbc://{USERNAME}:{PASSWORD}@{SERVER_NAME}/{DATABASE_NAME}?driver=ODBC+Driver+17+for+SQL+Server"
        engine = create_engine(conn_string)
        # Test connection
        with engine.connect() as conn:
            conn.execute(text("SELECT 1"))
        return engine
    except Exception as e:
        st.error(f"Failed to connect to SQL Server: {e}")
        return None

@st.cache_data
def get_database_schema(_engine):
    """Retrieve database schema information"""
    if not _engine:
        return ""
    
    try:
        inspector = inspect(_engine)
        schema_info = "Database Tables:\n\n"
        
        for table_name in inspector.get_table_names():
            schema_info += f"Table: {table_name}\n"
            columns = inspector.get_columns(table_name)
            for column in columns:
                schema_info += f"  - {column['name']} ({column['type']})\n"
            schema_info += "\n"
        
        return schema_info
    except Exception as e:
        st.error(f"Error retrieving schema: {e}")
        return ""

# Streamlit UI
st.set_page_config(page_title="AI SQL Chatbot", layout="wide")
st.title("🤖 AI SQL Chatbot")
st.markdown(f"Connected to SQL Server: **{SERVER_NAME}** | Database: **{DATABASE_NAME}**")

# Initialize OpenAI
openai_key = os.getenv("OPENAI_API_KEY")
if not openai_key:
    st.error("❌ OPENAI_API_KEY not found in .env file")
    st.stop()

# Connect to database
engine = get_sql_engine()
if not engine:
    st.error("❌ Could not connect to SQL Server. Check your connection credentials.")
    st.stop()

# Get database schema
schema_info = get_database_schema(engine)

if not schema_info:
    st.error("❌ Could not retrieve database schema.")
    st.stop()

# Display schema in sidebar
with st.sidebar:
    st.header("Database Schema")
    st.text_area("Available Tables and Columns", schema_info, height=400, disabled=True)

# Initialize LLM
model = ChatOpenAI(model="gpt-4", temperature=0, openai_api_key=openai_key)

# Create system prompt with schema context
system_prompt = f"""You are an expert SQL Server database assistant. Your job is to convert natural language questions into SQL Server T-SQL queries.

Important rules:
1. Only generate SELECT queries (no INSERT, UPDATE, DELETE, DROP, etc.)
2. Always use parameterized queries when possible
3. Include meaningful column aliases
4. Optimize for readability and performance
5. If the question is ambiguous, make reasonable assumptions and add a comment explaining them

Available Database Schema:
{schema_info}

Generate ONLY the SQL query, no explanations."""

prompt_template = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{input}")
])

# Chat interface
st.header("Ask Your Question")

user_question = st.text_area(
    "Enter your question in natural language:",
    placeholder="E.g., Show me all customers from New York with orders over $1000"
)

col1, col2 = st.columns([1, 1])

with col1:
    generate_btn = st.button("🔍 Generate SQL Query", key="generate")

with col2:
    clear_btn = st.button("🗑️ Clear", key="clear")

if clear_btn:
    st.session_state.user_question = ""
    st.rerun()

if generate_btn and user_question:
    try:
        st.divider()
        
        # Generate SQL query
        with st.spinner("Generating SQL query..."):
            chain = prompt_template | model | StrOutputParser()
            sql_query = chain.invoke({"input": user_question})
        
        # Display generated query
        st.subheader("Generated SQL Query:")
        st.code(sql_query, language="sql")
        
        # Execute query
        st.subheader("Query Results:")
        try:
            with st.spinner("Executing query..."):
                with engine.connect() as conn:
                    result = conn.execute(text(sql_query))
                    df = pd.DataFrame(result.fetchall(), columns=result.keys())
            
            if df.empty:
                st.info("Query executed successfully but returned no results.")
            else:
                st.dataframe(df, use_container_width=True)
                
                # Download option
                csv = df.to_csv(index=False)
                st.download_button(
                    label="📥 Download Results as CSV",
                    data=csv,
                    file_name="query_results.csv",
                    mime="text/csv"
                )
        
        except Exception as e:
            st.error(f"❌ Error executing query: {str(e)}")
            st.info("Please review the generated SQL query and try again.")
    
    except Exception as e:
        st.error(f"❌ Error generating SQL query: {str(e)}")