import streamlit as st
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import google.generativeai as genai
from langchain.vectorstores import FAISS
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate



def get_conversational_chain():

    prompt_template = """
    Answer the question as detailed as possible from the provided context, make sure to provide all the details, if the answer is not in
    provided context just say, "Sorry, I don't know about this. You can ask another question.", don't provide the wrong answer\n\n
    Context:\n {context}?\n
    Question: \n{question}\n

    Answer:
    """

    model = ChatGoogleGenerativeAI(model="gemini-pro",
                             temperature=0.3)

    prompt = PromptTemplate(template = prompt_template, input_variables = ["context", "question"])
    chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)

    return chain



def user_input(user_question):
    embeddings = GoogleGenerativeAIEmbeddings(model = "models/embedding-001")
    
    new_db = FAISS.load_local("./faiss_index", embeddings, allow_dangerous_deserialization=True)
    docs = new_db.similarity_search(user_question)

    chain = get_conversational_chain()

    
    response = chain(
        {"input_documents":docs, "question": user_question}
        , return_only_outputs=True)

    print(response)
    st.write("Reply: ", response["output_text"])





st.set_page_config("DIAT Rakshak")
st.title('Introducing "Rakshak"- DIAT Assistant Chatbot�')
st.info("Please check here- https://ai.google.dev/tutorials/web_quickstart , to get your GOOGLE GEMINI API KEY. This is mandatory to use this chatbot")

GOOGLE_API_KEY = st.text_input("Please enter your GOOGLE GEMINI API KEY", type="password")
#password = st.text_input("Enter a password", type="password")
os.environ['GOOGLE_API_KEY'] = GOOGLE_API_KEY
	
	
user_question = st.text_input("Hello! I am Rakshak, your DIAT assistant. Please ask your query regarding DIAT, Pune.")

if user_question:
    user_input(user_question)