import gradio as gr
from langchain_groq import ChatGroq
from langchain_community.llms import Ollama
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.document_loaders import TextLoader
from langchain_community.document_loaders.csv_loader import CSVLoader
from langchain_community.embeddings import OllamaEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains import create_retrieval_chain
from langchain_community.vectorstores import FAISS
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from dotenv import load_dotenv
import os
import time
from datetime import datetime

load_dotenv()

print("Function started at",datetime.now())
start=datetime.now()


os.environ["TOKENIZERS_PARALLELISM"] = "false"

groq_api_key="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"
loader1=TextLoader("Chat/ffreedomsample.txt")
loader2=TextLoader("Chat/lotr.txt")
loader3=CSVLoader("at_!/courseclean100.csv")
docs=loader3.load()
print("Document loaded")

# print("Documents loaded in ", (datetime.now()-start).seconds)

text_splitter=RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
final_documents=text_splitter.split_documents(docs)
print("Text has been split at", datetime.now())


embeddings = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1",)
# embeddings = HuggingFaceEmbeddings(model_name="Alibaba-NLP/gte-Qwen2-1.5B-instruct",)

vectors=FAISS.from_documents(final_documents, embeddings)


print("Vectorization done in", (datetihme.now()-start).seconds, "seconds")

# st.title("ChatGroq Ffreedom Llama3")

# llm=ChatGroq(groq_api_key=groq_api_key,model_name="llama3-8b-8192")
llm=Ollama(model="llama3.1")

prompt=ChatPromptTemplate.from_template(
"""
You are Kevai, you are a virtual assistant for a company Ffreedom, which offers various courses in the domain of personal finance, small business and farming, help users find the right course or solve their queries based on the provided context.
intoduce youself once at the beginning thereafter only introduce when asked
Answer the user query based on data you are taking in.
Answer within 100 words.
{context}
</context>
Question: {input}
"""
)
document_chain = create_stuff_documents_chain(llm, prompt)
retriever = vectors.as_retriever()
retrieval_chain = create_retrieval_chain(retriever, document_chain)

# prompt=st.text_input("Enter Query")

while True:
    t=input("Enter Query\n\n")
    if(t is not None):
        ret=datetime.now()
        response=retrieval_chain.invoke({"input":t,"prompt":prompt})
        print("\n")
        print(response['answer'], "\n")
        print("Response in", (datetime.now()-ret).seconds, "seconds\n\n")
    else:
        print("incorrect input")

gr.ChatInterface(
    chat_function,
    textbox=gr.Textbox(placeholder="Welcome aboard!", container=False, scale = 9),
    textbox=gr.Textbox(placeholder="how can i help you further", container=False, scale = 3)
    # additional_inputs=[gr.Textbox("You are helpful AI", label="System Prompt"),gr.Slider(500,4000, label="Max New Tokens"),gr.Slider(0,1, label="Temperature")]
    ).launch()