from langchain.utilities import SQLDatabase
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain_experimental.sql import SQLDatabaseChain
from dotenv import load_dotenv, find_dotenv

# Load environment variables
load_dotenv(find_dotenv())

# Set up the llm
llm = ChatOpenAI(temperature=0.0)

db = SQLDatabase.from_uri("postgresql://luminlab:Lumin#adapt2023@194.29.101.181:5433/luminlab_db")
ber_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)

if __name__ == "__main__":
    print("Running SQL chain")
    ber_chain.run('What is the average energy consumption of the building?')