app.py

import os
import openai
import streamlit as st
from sqlalchemy import create_engine
from langchain_community.agent_toolkits.sql.base import create_sql_agent
from langchain_community.utilities import SQLDatabase
from langchain.agents.agent_types import AgentType
from langchain_community.callbacks.streamlit import StreamlitCallbackHandler
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from langchain.chat_models import ChatOpenAI

# Ensure required dependencies are installed
try:
    import mysql.connector  # Required for MySQL connection
except ImportError:
    st.error("❌ Missing dependency: Install mysql-connector-python with `pip install mysql-connector-python`")
    st.stop()

# Set up Streamlit page
st.set_page_config(page_title="AI SQL Assistant", page_icon="🤖", layout="wide")
st.title("🤖 AI SQL Assistant")

# Load OpenAI API Key securely from environment variable
openai.api_key = os.getenv("OPENAI_API_KEY")
if not openai.api_key:
    st.error("❌ OpenAI API key is missing! Set it in your environment variables.")
    st.stop()

# Function to interact with OpenAI API
def get_chatgpt_response(prompt):
    try:
        client = openai.OpenAI()  # Initialize OpenAI client
        response = client.chat.completions.create(  # New API method
            model="gpt-3.5-turbo",  
            messages=[
                {"role": "system", "content": "You are an AI assistant that helps with SQL queries."},
                {"role": "user", "content": prompt},
            ],
        )
        return response.choices[0].message.content.strip()  # Updated response structure
    except Exception as e:
        return f"Error with OpenAI API: {e}"

# Sidebar for database configuration
with st.sidebar:
    st.header("⚙️ Database Configuration")
    selected_opt = st.radio("Choose an option", ["How to Use", "Connect to MySQL Database"])

    if selected_opt == "Connect to MySQL Database":
        mysql_host = st.text_input("MySQL Host", value="localhost")
        mysql_user = st.text_input("MySQL User")
        mysql_password = st.text_input("MySQL Password", type="password")
        mysql_db = st.text_input("MySQL Database")

        if st.button("Clear Chat History"):
            st.session_state["messages"] = []
    else:
        st.markdown("ℹ️ **Instructions:** Enter your MySQL database details to connect and run SQL queries.")

# Validate database credentials
if selected_opt == "Connect to MySQL Database":
    if not all([mysql_host, mysql_user, mysql_password, mysql_db]):
        st.warning("⚠️ Please enter all database details.")
        st.stop()

    # Create SQLAlchemy connection string
    db_uri = f"mysql+mysqlconnector://{mysql_user}:{mysql_password}@{mysql_host}/{mysql_db}"

    try:
        engine = create_engine(db_uri)
        db = SQLDatabase(engine)
        st.success("✅ Successfully connected to the MySQL database!")
    except Exception as e:
        st.error(f"❌ Database connection failed: {e}")
        st.stop()

    # Set up LLM model
    llm = ChatOpenAI(model="gpt-3.5-turbo", openai_api_key=openai.api_key)

    # Create SQL toolkit
    toolkit = SQLDatabaseToolkit(db=db, llm=llm)

    # Streamlit callback handler
    streamlit_callback = StreamlitCallbackHandler(st.container())

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you with SQL queries?"}]

# Display chat history
st.header("💬 Chat Interface")
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.write(msg["content"])

# User input
user_query = st.chat_input(placeholder="Ask anything about your database...")

if user_query:
    st.session_state.messages.append({"role": "user", "content": user_query})
    st.chat_message("user").write(user_query)

    with st.chat_message("assistant"):
        response_container = st.container()
        with response_container:
            try:
                response = get_chatgpt_response(user_query)
                response_container.markdown(response)
                st.session_state.messages.append({"role": "assistant", "content": response})
            except Exception as e:
                error_message = f"⚠️ Error: {e}"
                response_container.markdown(error_message)
                st.session_state.messages.append({"role": "assistant", "content": error_message})