import langchain
from langchain.document_loaders import PyPDFLoader, DirectoryLoader
from langchain.prompts import PromptTemplate
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.llms import CTransformers
from langchain.chains import RetrievalQA
from flask import Flask, request, render_template
from datetime import datetime
from flask import Flask, render_template, request, jsonify, session


app = Flask(__name__)


@app.route('/')
def home():

    return render_template('index.html')



DB_FAISS_PATH = 'vectorstore/db_faiss'

custom_prompt_template = """Given the following context and a question, generate an answer based on this context only.
    In the answer try to provide as much text as possible from "response" section in the source document context without making much changes.
    If the answer is not found in the context, kindly state "I don't know." Don't try to make up an answer.

    CONTEXT: {context}

    QUESTION: {question}


    """

def set_custom_prompt():
    """
    Prompt template for QA retrieval for each vectorstore
    """
    prompt = PromptTemplate(template=custom_prompt_template,
                            input_variables=['context', 'question'])
    return prompt

#Retrieval QA Chain
def retrieval_qa_chain(llm, prompt, db):
    qa_chain = RetrievalQA.from_chain_type(llm=llm,
                                       chain_type='stuff',
                                       retriever=db.as_retriever(search_kwargs={'k': 2}),
                                       return_source_documents=True,
                                       chain_type_kwargs={'prompt': prompt}
                                       )
    return qa_chain
  
#Loading the model
def load_llm():
    # Load the locally downloaded model here
    # llm = CTransformers(
    #     model = r"C:\Aiproject\Llama-2-7B-Chat-GGML\Llama-2-7B-Chat-GGML\llama-30b.ggmlv3.q8_0.bin",
    #     model_type="llama",
    #     max_new_tokens = 512,
    #     temperature = 0.5
    # )
    llm = CTransformers(model=r"TheBloke/Mistral-7B-Instruct-v0.1-GGUF",gpu_layers=100,config={'max_new_tokens': 128, 'temperature': 0.01})
    return llm

#QA Model Function
def qa_bot():
    embeddings = HuggingFaceEmbeddings(model_name=r"sentence-transformers/all-MiniLM-L6-v2",
                                       model_kwargs={'device': 'cpu'})
    db = FAISS.load_local(DB_FAISS_PATH, embeddings)
    llm = load_llm()
    qa_prompt = set_custom_prompt()
    qa = retrieval_qa_chain(llm, qa_prompt, db)

    return qa

#output function
def final_result(query):
    qa_result = qa_bot()
    response = qa_result({'query': query})
    return response

# Streamlit application
# st.title("Medical Bot")

# if "query" not in st.session_state:
#     st.text("Hi, Welcome to Medical Bot. What is your query?")
#     query = st.text_input("Enter your question")
#     st.button("Ask")

#     if query:
#         st.text("Searching for relevant documents...")
#         response = final_result(query)
#         st.text(f"Sources: {response['source_documents']}")
#         st.text(f"Answer: {response['result']}")
# else:
#     st.text("No query")




@app.route('/user_input', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        query = request.form.get('user_input')
        session_id= request.form.get('session_id')
        #query = request.form['query']
        start=datetime.now()
        frmtstart=start.strftime('%Y-%m-%d %H:%M:%S')
        print('started',frmtstart)
        user_input=query.lower()
        print(user_input)
        if query:
            #result = chain(question)['result']
            response = final_result(user_input)
            result=response['result']
            print(result)
            end=datetime.now()
            endtime=end.strftime('%Y-%m-%d %H:%M:%S')
            print('Ended',endtime)

      
        else:
            result = None
        #return render_template('index.html', output=result)
        return jsonify({'output': result})
    return render_template('index.html', result=None)



if __name__ == '__main__':
    app.run(host="0.0.0.0",port=8000,debug=False)