from langchain import OpenAI, SQLDatabase
from langchain_experimental.sql import SQLDatabaseChain
from langchain.memory import ConversationBufferMemory
from langchain.agents import (AgentType,
                              AgentExecutor,
                              create_react_agent, 
                              create_openai_functions_agent,
                              create_sql_agent)
from langchain.agents.initialize import initialize_agent
from langchain.tools import Tool
from langchain.prompts import PromptTemplate
# from langchain.retrievers import RetrievalModel
from langchain.indexes.vectorstore import VectorstoreIndexCreator
from langchain_community.document_loaders.text import TextLoader
from langchain_community.agent_toolkits import SQLDatabaseToolkit

# Document Loaders for PDFs and Docx
from langchain.document_loaders.pdf import PyPDFLoader
from langchain.document_loaders.word_document import Docx2txtLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores.chroma import Chroma
from langchain.chat_models.openai import ChatOpenAI
from langchain.chains import ConversationalRetrievalChain
from langchain.text_splitter import CharacterTextSplitter

import langchain
import psycopg2
import os
import constants
import constants_allen
os.environ["OPENAI_API_KEY"] = constants.OPENAI_API_KEY

import warnings
warnings.filterwarnings('ignore')


# Setup database connection
# Setup database
db = SQLDatabase.from_uri(
    f"postgresql+psycopg2://postgres:{constants.DBPASS}@localhost:5433/{constants.DB}",
)

# Setup conversation memory
memory = ConversationBufferMemory(memory_key="chat_history")

# Setup llm
llm = OpenAI(temperature=0, model_name="gpt-3.5-turbo")

# Create a list of documents from all our files in the ./data/docs folder
documents = []
for file in os.listdir("./data/docs"):
    if file.endswith(".pdf"):
        pdf_path = "./data/docs/" + file
        loader = PyPDFLoader(pdf_path)
        documents.extend(loader.load())
    elif file.endswith("./docx") or file.endswith(".doc"):
        doc_path = "./data/docs/" + file
        loader = Docx2txtLoader(doc_path)
        documents.extend(loader.load())
    elif file.endswith(".txt"):
        text_path = "./data/docs/" + file
        loader = TextLoader(text_path)
        documents.extend(loader.load())

# Split the documents into smaller chunks
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=10)
documents = text_splitter.split_documents(documents)

# Convert the document chunks to embedding and save them to the vector store
vectordb = Chroma.from_documents(documents, embedding=OpenAIEmbeddings(), persist_directory="./data/docs")
vectordb.persist()

# Setup prompt template
_DEFAULT_TEMPLATE = """Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer. Unless the user specifies in his question a specific number of examples he wishes to obtain, always limit your query to at most {top_k} results. You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for a the few relevant columns given the question.
Pay attention to use only the column names that you can see in the schema description. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.

Use the following format:

Question: Question here
SQLQuery: SQL Query to run
SQLResult: Result of the SQLQuery
Answer: Final answer here

Only use the following tables:

{table_info}

Question: {input}
"""

# TEMPLATE = '''"Given your input question, provide the answer from the most relevant source. You can choose between two sources:

# 1. Text File Source: This source contains Q&A data on data structures and algorithms.

# 2. Database Source: This source includes information from three tables:
#     - 'developers' table: Contains details about developers, such as full_name, email, phone, position, and department.
#     - 'tasks' table: Holds information about developers' tasks, including the task, completion status, due date, completion date, and priority.
#     - 'insurance_data' table: Contains information about PRU Life Insurance and USA Insurance, including questions and answers.

# Only query the content of the sources, not the metadata.
# You have access to the following tools:

# {tools}

# Use the following format:

# Question: the input question you must answer
# Thought: you should always think about what to do
# Action: the action to take, should be one of [{tool_names}]
# Action Input: the input to the action
# Observation: the result of the action
# ... (this Thought/Action/Action Input/Observation can repeat N times)
# Thought: I now know the final answer
# Final Answer: the final answer to the original input question

# Begin!

# Question: {input}
# Thought:{agent_scratchpad}'''

PROMPT = PromptTemplate.from_template(_DEFAULT_TEMPLATE)

# Setup db chain
db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)
toolkit = SQLDatabaseToolkit(db=db, llm=OpenAI(temperature=0, verbose=True, model_name="gpt-3.5-turbo"))

# Create Document Q&A chain
pdf_qa = ConversationalRetrievalChain.from_llm(
    OpenAI(temperate=0.7, model_name="gpt-3.5-turbo"),
    retriever=vectordb.as_retriever(search_kwargs={'k': 6}),
    return_source_documents=True,
    verbose=False
)

# Setup text data source
directory = "./data/text/final1.txt"
loader = TextLoader(directory)
index_creator = VectorstoreIndexCreator()
index = index_creator.from_loaders([loader])
retriever_text = index.vectorstore.as_retriever()

# Define the tools needed
tools = [
    Tool(name="MyDocs", func=pdf_qa, description="Retrieve information from documents."),
    Tool(name="MyDB", func=db_chain.run, description="Query the PostgreSQL database",),
    Tool(name="MyRetriever", func=retriever_text.get_relevant_documents, description="Retrieve documents from text file"),
]

agent = initialize_agent(
    tools, llm, memory=memory,
    agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
    verbose=True,
    handle_parsing_errors=True, 
    output_key="result",
    early_stopping_method="generate",
    max_iterations=3,
)

# agent = create_sql_agent(
#     llm=OpenAI(temperature=0),
#     toolkit=toolkit,
#     verbose=True,
#     agent_type=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
# )

# agent = create_react_agent(
#     llm=llm,
#     tools=tools,
#     prompt=PROMPT
# )

# agent_executor = AgentExecutor(agent=agent, tools=tools, memory=memory,
#                                verbose=True, handle_parsing_errors=True)

yellow = "\033[0;33m"
green = "\033[0;32m"
white = "\033[0;39m"

def chat():
    while True:
        user_input = input("You: ") # Get user input

        # Use the agent to generate a response
        try:
            bot_response = agent.run(user_input)
            # bot_response = agent_executor.invoke({"input": user_input})
        except ValueError as e:
            bot_response = str(e)
            if not bot_response.startswith("Could not parse LLM output: `"):
                raise e
            bot_response = bot_response.removeprefix("Could not parse LLM output: `").removesuffix("`")

        print("Chatbot: " + bot_response)

        if user_input.lower() == 'exit':
            break

if __name__ == "__main__":
    chat()
langchain.text_splitter.CharacterTextSplitter langchain_experimental.sql.SQLDatabaseChain langchain.agents.initialize.initialize_agent langchain.tools.Tool langchain.document_loaders.pdf.PyPDFLoader langchain.memory.ConversationBufferMemory langchain_community.document_loaders.text.TextLoader langchain.document_loaders.word_document.Docx2txtLoader langchain.indexes.vectorstore.VectorstoreIndexCreator langchain.prompts.PromptTemplate.from_template langchain.SQLDatabase.from_uri langchain.OpenAI langchain.embeddings.openai.OpenAIEmbeddings
[((1372, 1405), 'warnings.filterwarnings', 'warnings.filterwarnings', (['"""ignore"""'], {}), "('ignore')\n", (1395, 1405), False, 'import warnings\n'), ((1458, 1572), 'langchain.SQLDatabase.from_uri', 'SQLDatabase.from_uri', (['f"""postgresql+psycopg2://postgres:{constants.DBPASS}@localhost:5433/{constants.DB}"""'], {}), "(\n    f'postgresql+psycopg2://postgres:{constants.DBPASS}@localhost:5433/{constants.DB}'\n    )\n", (1478, 1572), False, 'from langchain import OpenAI, SQLDatabase\n'), ((1608, 1659), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""'}), "(memory_key='chat_history')\n", (1632, 1659), False, 'from langchain.memory import ConversationBufferMemory\n'), ((1679, 1728), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0)', 'model_name': '"""gpt-3.5-turbo"""'}), "(temperature=0, model_name='gpt-3.5-turbo')\n", (1685, 1728), False, 'from langchain import OpenAI, SQLDatabase\n'), ((1831, 1856), 'os.listdir', 'os.listdir', (['"""./data/docs"""'], {}), "('./data/docs')\n", (1841, 1856), False, 'import os\n'), ((2402, 2458), 'langchain.text_splitter.CharacterTextSplitter', 'CharacterTextSplitter', ([], {'chunk_size': '(1000)', 'chunk_overlap': '(10)'}), '(chunk_size=1000, chunk_overlap=10)\n', (2423, 2458), False, 'from langchain.text_splitter import CharacterTextSplitter\n'), ((5005, 5052), 'langchain.prompts.PromptTemplate.from_template', 'PromptTemplate.from_template', (['_DEFAULT_TEMPLATE'], {}), '(_DEFAULT_TEMPLATE)\n', (5033, 5052), False, 'from langchain.prompts import PromptTemplate\n'), ((5082, 5134), 'langchain_experimental.sql.SQLDatabaseChain', 'SQLDatabaseChain', ([], {'llm': 'llm', 'database': 'db', 'verbose': '(True)'}), '(llm=llm, database=db, verbose=True)\n', (5098, 5134), False, 'from langchain_experimental.sql import SQLDatabaseChain\n'), ((5559, 5580), 'langchain_community.document_loaders.text.TextLoader', 'TextLoader', (['directory'], {}), '(directory)\n', (5569, 5580), False, 'from langchain_community.document_loaders.text import TextLoader\n'), ((5597, 5622), 'langchain.indexes.vectorstore.VectorstoreIndexCreator', 'VectorstoreIndexCreator', ([], {}), '()\n', (5620, 5622), False, 'from langchain.indexes.vectorstore import VectorstoreIndexCreator\n'), ((6067, 6289), 'langchain.agents.initialize.initialize_agent', 'initialize_agent', (['tools', 'llm'], {'memory': 'memory', 'agent': 'AgentType.CONVERSATIONAL_REACT_DESCRIPTION', 'verbose': '(True)', 'handle_parsing_errors': '(True)', 'output_key': '"""result"""', 'early_stopping_method': '"""generate"""', 'max_iterations': '(3)'}), "(tools, llm, memory=memory, agent=AgentType.\n    CONVERSATIONAL_REACT_DESCRIPTION, verbose=True, handle_parsing_errors=\n    True, output_key='result', early_stopping_method='generate',\n    max_iterations=3)\n", (6083, 6289), False, 'from langchain.agents.initialize import initialize_agent\n'), ((5321, 5370), 'langchain.OpenAI', 'OpenAI', ([], {'temperate': '(0.7)', 'model_name': '"""gpt-3.5-turbo"""'}), "(temperate=0.7, model_name='gpt-3.5-turbo')\n", (5327, 5370), False, 'from langchain import OpenAI, SQLDatabase\n'), ((5759, 5848), 'langchain.tools.Tool', 'Tool', ([], {'name': '"""MyDocs"""', 'func': 'pdf_qa', 'description': '"""Retrieve information from documents."""'}), "(name='MyDocs', func=pdf_qa, description=\n    'Retrieve information from documents.')\n", (5763, 5848), False, 'from langchain.tools import Tool\n'), ((5849, 5935), 'langchain.tools.Tool', 'Tool', ([], {'name': '"""MyDB"""', 'func': 'db_chain.run', 'description': '"""Query the PostgreSQL database"""'}), "(name='MyDB', func=db_chain.run, description=\n    'Query the PostgreSQL database')\n", (5853, 5935), False, 'from langchain.tools import Tool\n'), ((5937, 6058), 'langchain.tools.Tool', 'Tool', ([], {'name': '"""MyRetriever"""', 'func': 'retriever_text.get_relevant_documents', 'description': '"""Retrieve documents from text file"""'}), "(name='MyRetriever', func=retriever_text.get_relevant_documents,\n    description='Retrieve documents from text file')\n", (5941, 6058), False, 'from langchain.tools import Tool\n'), ((1946, 1967), 'langchain.document_loaders.pdf.PyPDFLoader', 'PyPDFLoader', (['pdf_path'], {}), '(pdf_path)\n', (1957, 1967), False, 'from langchain.document_loaders.pdf import PyPDFLoader\n'), ((2644, 2662), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (2660, 2662), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((5175, 5238), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0)', 'verbose': '(True)', 'model_name': '"""gpt-3.5-turbo"""'}), "(temperature=0, verbose=True, model_name='gpt-3.5-turbo')\n", (5181, 5238), False, 'from langchain import OpenAI, SQLDatabase\n'), ((2125, 2149), 'langchain.document_loaders.word_document.Docx2txtLoader', 'Docx2txtLoader', (['doc_path'], {}), '(doc_path)\n', (2139, 2149), False, 'from langchain.document_loaders.word_document import Docx2txtLoader\n'), ((2281, 2302), 'langchain_community.document_loaders.text.TextLoader', 'TextLoader', (['text_path'], {}), '(text_path)\n', (2291, 2302), False, 'from langchain_community.document_loaders.text import TextLoader\n')]