main.py

import streamlit as st
import pymysql
from sqlalchemy import create_engine
from dotenv import load_dotenv
from langchain_openai import AzureChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import create_sql_agent
from langchain.agents import AgentType
import os

# Load .env file
load_dotenv()

st.set_page_config(page_title="SQL AI Agent", page_icon="🤖", layout="centered")
st.title("🤖 AI SQL Agent with Azure OpenAI + MySQL")

# --- Session state variables ---
if "databases" not in st.session_state:
    st.session_state.databases = []
if "db_uri" not in st.session_state:
    st.session_state.db_uri = None

# Step 1: Get MySQL connection details
st.subheader("1️⃣ Enter MySQL Connection Details")
host = st.text_input("Host", value="127.0.0.1")
port = st.number_input("Port", value=3306, step=1)
user = st.text_input("Username", value="root")
password = st.text_input("Password", type="password")

if st.button("Connect to MySQL"):
    try:
        conn = pymysql.connect(
            host=host,
            port=port,
            user=user,
            password=password
        )
        with conn.cursor() as cursor:
            cursor.execute("SHOW DATABASES")
            st.session_state.databases = [db[0] for db in cursor.fetchall()]
        conn.close()

        st.success("✅ Connection successful. Select a database below.")
    except Exception as e:
        st.error(f"Connection failed: {e}")

# Step 2: Select database
if st.session_state.databases:
    selected_db = st.selectbox("Select Database", st.session_state.databases)
    if st.button("Use This Database"):
        st.session_state.db_uri = f"mysql+pymysql://{user}:{password}@{host}:{port}/{selected_db}"
        st.success(f"Using database: {selected_db}")

# Step 3: Start AI SQL Agent
if st.session_state.db_uri:
    try:
        llm = AzureChatOpenAI(
            deployment_name=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
            openai_api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
            openai_api_key=os.getenv("AZURE_OPENAI_API_KEY"),
            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
            temperature=0
        )

        db = SQLDatabase.from_uri(st.session_state.db_uri)

        agent_executor = create_sql_agent(
            llm=llm,
            db=db,
            agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            verbose=True
        )

        query = st.text_area("💬 Ask your question in English", height=100)
        if st.button("Run Query"):
            with st.spinner("Running query..."):
                try:
                    result = agent_executor.run(query)
                    st.write("### Result")
                    st.write(result)
                except Exception as e:
                    st.error(f"Error: {e}")

    except Exception as e:
        st.error(f"Agent setup failed: {e}")