# RUN THIS CELL FIRST!
#!pip install -q langchain openai tiktoken faiss-cpu streamlit InstructorEmbedding sentence-transformers gradio -q

from InstructorEmbedding import INSTRUCTOR
import streamlit as st
import os

from langchain.embeddings import HuggingFaceInstructEmbeddings
from langchain import LLMChain
from langchain.memory import ConversationBufferMemory
from langchain.vectorstores.faiss import FAISS
from langchain.prompts import PromptTemplate
import os
from googletrans import Translator
from langchain.chat_models import ChatOpenAI
import gradio as gr
import openai


OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
HUGGINGFACEHUB_API_TOKEN=os.environ["HUGGINGFACEHUB_API_TOKEN"]

translator = Translator()

def make_context(docs):
  context = ""
  for doc in docs:
    doc = doc.page_content +  "\n\nSource: " + doc.metadata['source']
    context = context + doc + "\n\n"
  return context

#simple function to detect and translate text
def detect_and_translate(text,answer):
    result_lang = translator.detect(text)
    print(result_lang.lang)
    if result_lang.lang == "en":
        translate_text = answer
    else:
        translate_text = translator.translate(answer,src='auto', dest=result_lang.lang).text
    return translate_text

def detect_and_translate_ques(text):
    result_lang = translator.detect(text)
    print(result_lang.lang)
    if result_lang.lang == "en":
        translate_text = text
    else:
        translate_text = translator.translate(text,src='auto', dest="en").text
    return translate_text


instructor_embeddings = HuggingFaceInstructEmbeddings(model_name="hkunlp/instructor-xl")
new_vectorstore = FAISS.load_local("./faiss_docs_xl_500_80_index", instructor_embeddings)


with gr.Blocks() as demo:
    
    chatbot = gr.Chatbot(height=300)
    msg = gr.Textbox(label="Ask Freddy")
    clear = gr.Button("Clear")
    chat_history = []
    def user(user_message, history):

        question=detect_and_translate_ques(user_message)
        print(question)        
        mychosendocs=new_vectorstore.similarity_search(query=question)
        context = make_context(mychosendocs)
        memory = ConversationBufferMemory(k=10,memory_key="chat_history")
        template = """
        your job is to answer the questions asked by the users. Create a final answer with references ("SOURCES").
        If the answer is not in the context, then say that you dont know.
        Source of the context is written at the end of the context.
        At the end of your answer write the source of the context in the following way: \n\nSource: (source)
        Chat history is also provided to you.
        Context: {context}
        ---

        Chat History: {chat_history}
        Question: {question}
        Answer: Let's think step by step and give best answer possible. Use points when needed.
        """

        llm = ChatOpenAI(model='gpt-3.5-turbo',temperature = 0, openai_api_key= OPENAI_API_KEY)

        prompt = PromptTemplate(template=template, input_variables=["context", "question", "chat_history"]).partial(context=context)

        llm_chain = LLMChain(prompt=prompt, llm=llm, verbose=False,memory=memory)

        response = llm_chain.run(question)
        if "Question:" in response:
          respone_new="Sorry, I don't know the answer! Can you please rephrase your question?"
        else:
          respone_new=response
        
        my_correct_answer=detect_and_translate(user_message,respone_new)        
        chat_history.append((user_message, my_correct_answer))
        
        return gr.update(value=""), chat_history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False)
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch(debug=False,share=False)