from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.llms import CTransformers
from langchain.memory import ConversationBufferMemory
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import json

DB_FAISS_PATH = 'vectorstore.faiss'

custom_prompt_template =  """
Answer the question based on the context below. Keep the answer short and concise. Respond "Unsure about answer" if not sure about the answer. If you find context is unrelevant for the question then ignore context and answer question without use of context.

Context: {context}
Current conversation:{history}
Human: {question}

Provide the resolution steps directly, it may be a new customer who is having same question, to overcome the issues also provide the instructions to resolve the issue
Helpful resolution:
"""

def set_custom_prompt():
    prompt = PromptTemplate(template=custom_prompt_template,
                            input_variables=['history','context', 'question'])
    return prompt

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}),
                                       verbose=True,
                                       return_source_documents=True,
                                       chain_type_kwargs={'verbose': True,
                                                          'prompt': prompt,
                                                          'memory': ConversationBufferMemory(
                                                              memory_key="history",
                                                              input_key="question",
                                                              k=3
                                                          )}
                                       )
    return qa_chain

def qa_bot():
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2",
                                       model_kwargs={'device': 'cpu'})
    db = FAISS.load_local(DB_FAISS_PATH, embeddings)
    llm = CTransformers(model="mistral")
    qa_prompt = set_custom_prompt()
    qa = retrieval_qa_chain(llm, qa_prompt, db)
    return qa