from sqlalchemy import create_engine
from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits.sql.base import create_sql_agent  # Updated Import
import streamlit as st
import os
from dotenv import load_dotenv

load_dotenv()

def sql_agent(prompt):
    # Set OpenAI API key
    OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

    DATABASE_URL = "sqlite:///data/main.db"  # Change to your database URL
    engine = create_engine(DATABASE_URL)
    db = SQLDatabase(engine)

    # Initialize OpenAI model
    llm = ChatOpenAI(temperature=0, model_name="gpt-4o")

    # Create SQL Agent
    agent_executor = create_sql_agent(
        llm=llm,
        db=db,
        agent_type="openai-tools",  # Updated agent type
        verbose=True
    )

    # Example query
    # query = "How many Senior Engineers are in the company?"
    response = agent_executor.invoke({"input": prompt})['output']
    
    return response