import os
from langchain import HuggingFaceHub
from langchain import PromptTemplate, LLMChain
from langchain.sql_database import SQLDatabase
from langchain.chat_models import ChatOpenAI
from langchain.agents import AgentExecutor
from langchain.agents.agent_types import AgentType
from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
import gradio as gr
from gradio.components import Textbox
from gradio.themes import Soft

database_user = os.environ.get("database_user")
database_password = os.environ.get("database_password")
database_server = os.environ.get("database_server")
database_db = os.environ.get("database_db")

connection_string = f"mssql+pymssql://{database_user}:{database_password}@{database_server}.database.windows.net:1433/{database_db}"
db = SQLDatabase.from_uri(connection_string)

API_KEY = os.environ.get("API_KEY_NEW")

llm = ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key=API_KEY, temperature=0)

toolkit = SQLDatabaseToolkit(db=db, llm=llm)

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

def langchain_bot(prompt):
    if prompt:
        final_prompt = prompt + ' Note: After getting the final answer correctly, set the final answer value to the same as the last "Thought", for example: Thought: The result is X, Final Answer: The result is X'
        text = agent_executor.run(final_prompt)
        return text
    return "Type a prompt to interact with Adventure Works Data Warehouse !"

iface = gr.Interface(fn=langchain_bot, inputs=Textbox(lines=5, placeholder="Type your prompt here..."),
                     outputs=Textbox(), title="Adventure Works Data Warehouse AI Assistant 🤖",
                     theme=Soft(),
                     allow_flagging='never')

# Launch the interface
iface.launch(inline=False)